Flash Inference: Near Linear Time Inference for Long Convolution Sequence Models and Beyond
Authors: Costin-Andrei Oncescu
Sanket Purandare
Stratos Idreos
Sham M. Kakade
Venue: ICLR 2025 Poster (2025)
Keywords: computational efficiency
inference
sequence generative models
long convolution sequence models
hyena
Primary area: generative models
TL;DR: We introduce a general method for speeding up exact inference in some classes of generative sequence models

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 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.

Published as a conference paper at ICLR 2025
FLASH INFERENCE: NEAR LINEAR TIME INFERENCE FOR LONG CONVOLUTION SEQUENCE MODELS AND BEYOND
Costin-Andrei Oncescu, Sanket Purandare, Stratos Idreos & Sham Kakade Harvard University fconcescu, sanketpurandare @g. hww.ele........eeeeoo harvard. edu. {stratos, sham) harvard. edu
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 computa- tional issue. Some them, including long convolution sequence models (LCSMs), address this issue at training time but remain quadratic during inference. We pro- pose method for speeding up LCSMs' exact inference 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 a tiling which helps decrease memory movement and share computation. has the added benefit allowing for almost complete paral- lelization across layers of the position-mixio part of the architecture. Empirically, we provide a proof of concept implementation for Hyena, which gets up to 7.8x end-to-end improvement over standard inference by improving up 110x within the position-mixing part.
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. 2017). While these models have great quality, it comes at computation cost which scales quadratically 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 & Dao 2023 Poli et al., 2023] Fu et al. 2024). Although most of these works have improved computational efficiency for training, some still scale quadratically in sequence length when comes to inference, thus not improving asymptotically over transformers.
In this work, propose framework for optimizing inference efficiency for general class of such models. As case study, which inspired the method, we focus on long convolution sequence models (LCSMs) (Poli al. 2023 Fu al. 2022) Romero et al.] 22021b Li et al. 2022 Karami Ghodsi 2024, Fu et al. [2023a] Romero et al. 2021a). 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, the building block of the architecture is that convolving the input sequence with a sequence-length long, (potentially underparameterizedd filter. If we let L be the sequence length (c.g. number tokens in the case of LLMs), then a naive implementation of employ FFT to bring that down
convolution during training would take Q(L2) FLOPs, but to OLL log L). During inference. 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 2(L2) this is the apparent cost of moving from a static input to dynamic one.
turns out that in fact, at least in the case of "dynamic FFT", one can obtain O(L log L) time complexity by using van der Hoeven (1997)'s result relaxed polynomial convolution. Crucially,

---

Published as a conference paper at ICLR 2025
this result works by direct reduction several applications of FFT. This allows phrase more general framework for turning training-efficiem. architectures to inference-efffcient ones.
8 =1
2
3
4
5
6
7
8
y/1
4
21
2
Di
2
4
3
9[2.2]
Z2
M/
M]jii] 77
1
3
911,2,
5
#[9.4]
M[a.a]
5/4
4
4.
4.
3
4
F[4.4]
24
5
2
3[5.55]
V/G
J
contribution of V[1.4] to 2[5.8]
2[6.6]
20
3
4.
7
1
9[6.0]
27
D5
#[7.8]
M)?.7]
Ve
5
944.91
1
8
4
1=1
3
4
8 1=1
2
3
4
6
7
8
Figure 1: Cell (i, j) corresponds to the contribution of Vi to 2j. To compute 2j. all its line of contributions should've been accounted for. Because of the autoregressive nature of inference, one only has access to Y4 after 2 has been computed. (Left Top) represents the standard lazy approach, (Left Bottom) represents the eager approach, and (Right) represents our suggested method.
Our main contributions are:
(1) Proposing the first LCSMs. This is in contrast to previous work of (Massaroli et al. 2024) who performs approximate inference.
OL log" L). exact inference algorithm for
(2) Identifying the main features of LCSMs that allow for faster inference to propose a more general framework.
(3) Highlighting and exploiting the potential for dealing with most workload of different layers in parallel, rather than having wait for all computation of layer to finish before moving on to subsequent layers.
(4) Saving up on data movement thanks the tiling approach used. We show how save factors of in several places, including activation storage, when the convolutional filters are data-independent. However, our method extends to any causal, data-dependent filters.
(5) Our framework provides O( log² L) inference algorithm which, when empirically, yields up to 7.8x end-to-end time-efficiency improvement; and up to 110x on the long convolution part.
SETTING AND RELATED WORK
Recent efforts have focused on designing more efficient alternatives to transformers (Tay et al. 2022), with state space models (SSMs) (u et al. 202 Gu & Dao 2023) and convolution based models (Fu et al. 2022] Poli et al. 2023) emerging as notable contenders. In this work, we focus on the latter simply because leverage the convoluional view directly. Nevertheless, it is worth
2

---

Published as a conference paper at ICLR 2025
mentioning that there is one-to-one mapping between linear-time invariant (LTT) SSMs and LCSMs and consequently, our findings apply to LTI SSMs as well (since they represent the same models). While LCSMs inherently assume sequence-long convolutions, they can correspond to LTI SSMs of varying state dimensions (up to filter length). Using the recurrent formulation SSMs for inference incurs costs that scale with the state dimension, which become prohibitive for larger dimensions. In contrast, our method exclusively leverages the convolutional perspective and has no dependency on filter properties. An extended contextual discussion available in Appendix A
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 wise, feature-mixing mixing modules, partially following AArora et al. (2023)'s notation. Put formally, let block' .... block M IR + RP be feature-mixing modules and mixer' mixer IR 1x= + R be position-mixing modules (we use superscripts throughout to denote layers). One can think of the blocks as being combinations of MIPs, layernorms, skip connections and other modules of similar purpose and of mixers as playing the role of multi-head attention (MHA) in 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 Lmoz). We { M and position 1 S i L x represents the input embeddings (L
define the activations a" ...... aM RRXD - R at level as at (x)); - block" (mixer'(a" (a-=) where (x) D-dimensional tokens) and the model output at position i is read from the last layer a (x);
We focus on autoregressive sequence models where for an input sequence Fi,... Tp. one generates the (p + 1)th token by sampling from distribution given by a (z), (usually via a linear map). For this to be well-defined, we constrain (by construction of mixers) the models to be causal, that is, a (z); is a function of only 2[[, or, equivalently, mixer (y); is a function of V1....)
We use @ to denote the element-wise (Hadamard) product: for two matrices A, B E RRXM A B E N is given by (A (AOB),, A B.j. We call function f X + X' associative, f( 22 ................................ xp) = f(f(z1,... 2;), f(=++),.. kp)) for any 1 S i < and I Tp E X. Finally, we define the time complexity of an algorithm as the number of soating-point operations (FLOPs) executes this is notion independent of hardware.
2.2 SELF-ATTENTION
Self-attention is the building block transformers (Vaswani et al. 2017). ). In its simplest form, implements a mixer by using three projection matrices Q, K. V E R DDD to obtain three sequences of vectors q, k. U E R.x by projecting the input E Rix via q = yQ, A = yK, v = yV these called queries, keys and values, respectively. The (causal) self- self-attention operator attention RIxD + L.. 1 is then given by:
allention(y); A mixer(y); = El_ Uj ((
(Up1a)) softmax(ki..
(I)
E elo,.k.)
which represents, at position j, the average of the previous value vectors U|1,]] exponentially-weighted by how well the values' corresponding keys match jth query. Put differently, the bigger (q:, kz)) is. the more jth output attends to 1th 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 1j E Dx11). ). In the transformer architecture this is how one "head" operates and each embedding is normally split into several such heads this called multihead attention (MHA). If we take the causality away (that is, the i < J) for simplicity, one could think of attention as being mixerr()) = softmax(gk: )r where softmax is computed along the rows of what is called the attention matrix: qk.
2.3 LONG CONVOLUTION SEQUENCE MODELS (LCSMS)
LCSMs (Poli al. 2023
2024) work by creating SISO (single-input- single-output) primitive map input sequence y R RL via 2, 2 Ei= We Pt- where
Li et al.
2022 Gaido et al.
3

---

Published as a conference paper at ICLR 2025
P is a (possibly infinite, and of length at least L) convolution filter which is often parameterized by smaller latent space θ: = f()) where E O dim θ < is learnt. These convolution primitives operate
independent axes of the hidden dimensions to create positional mixer: mixer(y); - Ei=1 Yi @ Pt- TR D. This assumes the filters to be independent, but shared ones possible as well (as is the case for multi-head Hyena (Massaroli et al. 2024)).
For example, the Hyena architecture (Poli al. 2023) maps back our setup when mixers are defined as above and the block functions, depending on the layer, are either MLPs or gates - that is, wiss multiplications with projection of activations at same position, but lower level. We do not focus on the details of the blocks since they all involve some D x D matrix-vector multipiication that is performed once for every layer and position S L and thus scale as (LLD²) per layer- that is, linearly in context length.
2.3.1 THE INFERENCE PROBLEM
Whereas during training, the convolutions can be performed in OoL L) time via FFT as the whole of y is known beforehand, during inference this non-trivial. Suppose there is prompt of size P and we want to generate I, P more tokens. The bottleneck problem can be understood as computing, for every layer and every dimension, the below:
21
U1 Pt-i
(2)
for all 1 L. To pre-fill the first P values inputs are known, one can perform FFT they would normally do during train-time (incurring an O(Plog P) time cost), but past psh position it is important to note that yi is not available before computing 2;-1 this is where static FFT breaks. Since dealing with the prompt can be done easily (Massaroli et al. 2024 Lemma 2.1) 1) by essentially filling in all contributions of Vp...r to 2|1...] and then lorgetting the prompt ever existed, we henceforth assume P P=0 O.
at all levels of activations, since the first P
2.3.2 PREVIOUS WORK ON EFFICIENT INFERENCE
Speeding up Equation 2 from the naive S(LZ) is the object of study of Massaroli al. (0024): they optimize the process by finding a low-dimensional LTI SSM whose equivalent convolution filter closely resembles p. thus getting an RNN-like problem to simulate. If the learnt SSM has size D', then this yields an (I..DD') algorithm for generating T. tokens. This has the significant added benefit of not needing to store the activations (or even inputs) thus far which makes memory efficient and practically time-efficient.
The downside, however, is that by definition this will only an approximation of the learned filter. More importantly, this approximation represents a projection to a potentially different space of models, making the end-to-end procedure practically different training procedure for low- dimensional LTI SSM model. Furthermore, this approach assumes the filter p is data- fata-independent this is needed in order to undergo an expensive distillation procedure as precomputation step. This can be qualitatively limiting shown in Arora al. (2023). While do store all activations, our proposed framework exactly simulates the architecture and data-dependent filters can also be accommodated.
FAST LCSM INFERENCE
In this section we describe our method for the case of LCSMs. That is, we still work with models as described in Section [.1 but we further assume the mixers to be convolution-based, as described 2.3 Hence, for every mixer E (mixer,.... miixxrr }, there exists filter pe E RL.XD such that
mixer = Σ Of
We assume here, for simplicity, that p is data-independent and part of the model, as is the case for most existing architectures. However, discuss in Appendix C how our method can be extended
4

---

Published as a conference paper at ICLR 2025
to data-dependent filters (by paying factor of 2). We also assume L = 2P for some integer P for simplicity one can always round L up to the closest power of 2 without changing by more than a factor of 2 and thus keeping the asymptoties of all our analysis unchanged.
3.1 THE PROPOSED ALGORITIMM
We derive inspiration from the the work of van der Hoeven (997) regarding relaxed polynomial interpolation they propose fix to dealing with dynamic structure of Eq. to achieve an overall O(Llog L) time complexity.
3.1.1 RELAXED POLYNOMIAL INTERPOLATION
Consider Eq
21 Vi Pt
The problem of relaxed polynomial interpolation can be thought of having to compute 2, for every 1 & t S L with the further constraint that Y/ and Pr are only made available once 2 has been outputted. While this general setting is covered in Appendix C here focus the case that all of is known ahead of time and only y is incrementally revealed. Whereas one could simply use the more general approach of (van der Hoeven, 1997). we can further take advantage of p being known ahead of time to get a twice-faster and slightly simpler algorithm. We present this both for brevity and because existing architectures have data-independent p's so this is the more efficient choice for them (and thus, for practical purposes). However, (Arora et al. [023] Karami & Ghodsi (024) show that data-dependent filters can improve quality but finding good causal architecture remains open.
The main approaches to relaxed polynomial interpolation, as identified in (Van der Hoeven 2002). are:
The lazy approach: for every E {1... LJ. one computes =e by simply applying the formula. It takes (()) FLOPs to do and the overall complexity is O(L²). This approach can be thought of as the naive approach which only performs work when it is strictly needed.
the opposite end, there the eager (or zealous) approach: soon as a new value of y or p becomes available, one accounts for all its contributions to Since we assumed the entire p known beforehand, this translates to the following:
After computing #t-1, V/ becomes available, so for every t S S L, we increase & by Vi Pi--- 24 is now fully computed so we proceed to the next iteration.
Hence, the eager approach be thought of performing work performed. Here, gt iteration takes O(L t) so we still have an overall complexity of (
it can be
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 yi Pÿ) in such way that fast convolution methods (such as FFT) can be employed for a speedup.
First, let us define the sort of contribution grouping that we consider:
Definition. For any < < s Sr' o < L we let T(y, [1, p. [t', r')) represent the aggregated contribution of all yut) to all of zy That is, for all l' S S r', we have:
T(y, U1, r.
y Pt-i
The employment of FFT for speeding up the computation rly [t. rl, p. ('rr can then be achieved following the next lemma:
and L2 r' -1' +' of and z, respectively, There exists an FFT-based algorithm ranning in O(Li + L2) space and ((L + L2) log(L + Lz)) time complexity that, given access yu.r) computes all the elements of (v. 11, r], P, (('rr'D) that is, all the aggregated contributions of yp.r to all of s.
Lemma 1. Let 1
S
< r' < L represent ranges of lengths Li -
5

---

Published as a conference paper at ICLR 2025
What this lemma says is that we can compute efficiently the contribution of range of inputs to range of outputs, where the computational cost is asymptotically given by the larger of the ranges (this is because [LL + L2) because in order to account for all these contributions at once, we need to have all of V[trr] available which means that all of F[1,r-1] was computed, by that time should have already accounted for the contributions of Virr1] to all of 211,r-11; hence, there is no point to (re)computing any of these.
O (max(LL L2)). Furthermore, observe that we assumed r S L': this
The lazy, eager, and our version of relaxed approach are depicted in Figure 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 to one element of 2 in the lazy approach and that of singe element of to long range of 2 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 I l' and one needs to ensure that the whole line of contribution to some 2n has been dealt with by the time return the value of 2, this is indeed the case for our tiling.
Algorithm 1 Fast Relaxed Polynomial Interpolation
Require: filter p R- and y 1: Initialize 2 E _R_ with zeros
2: for
to L 1 do
3: U < maximum power of 2 that divides # the side of the 8th gray tile
4: 2, += V/ * Po finalize 2; by accounting for newly available ys - red cell
5: # account for the contribution of U/--U+1a] to 2[441,407] gray tile
6: 2/+4 1. += Tv, [ U + 1, [ + 1, UT) # I is defined as in Eq
7: return 2g # 24 has been computed
8: unlock
# ys++1 becomes available
9: end for
10:
Vr. PD
11: return z1 z7 has been computed
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 Yt before zt-1 has been computed. This is shown in Algorithm [ return marks a value of z having been computed and wnlock represents a value of y becoming available. At the beginning of iteration it always holds that
2 E}= Bj Pi-j. We start the iteration by completing this with the freshly unlocked y; Po at line 4 this corresponds to the red tiles of Figure | Following this step 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(+-0+11 1.4] 2: (.+UU where U is the side of the tile, namely the largest power of dividing i. Finally we return as being in its final form and get access to Vi+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 I. 2P, Algorithm D perforrs 2P-1-9 calls to on ranges of length 29 for every a E (0... ... P -1.. J. If we base T's 's implementation on Lemma 7 this entails a total time complexity of O(L L). Since the calls to are sequential and outputs are not stored, the peak memory usage is given by the largest call and remains (L)).
3.1.2 PERFORMING LCSM INFERENCE VIA FAST RELAXED POLYNOMIAL INTERPOLATION
We can now use Algorithm to get an efficient inference algorithm for LCSMs, but to do so we need to be careful about dependeeend. First, for brevity, in our previous notation of activations, disregard the dependency on input x that is, use of instead of a (z)); ; since we are only going to work with one input x (the one we are generating). Furthermore, denote the intermediate mixer

---

Published as a conference paper at ICLR 2025
computation by b RxxxxD as follows:
of - mixer aa= 1a=
Σ k=
a
2-1 o pi
(4)
of block (b.)
(5)
for any 1 S 8 S M, where of is the embedding of 7th token. Whenever one has generated the tokens 2|1,4| up to position or, equivalently, Li..j the whole of [0.M]
[1,-11] will have been computed
(since a[ is needed sample -'(, Caching these up the equivalent of KV-cache in
transformers and is simply the natural way to repeat work when autoregressively sampling. To
generate the next token 21+1 = 4.+1. we want to fill the values of a. [ To do so, as per Eq. 4
we need to compute of which is practically, for each level e, and within each dimension < D.
relaxed interpolation between df (playing the role of y) and pf. Further applying block to b one gets a.. Thus, can pipeline several instances of Algorithm sequentially across layers.
The resulting algorithm is illustrated in Algorithm 2 The outer loop (line 3 iterates through positions each iteration, one computes the activations at position across all layers, does some (limited) eager work and samples next token. To do so, the inner loop (line iterates through the several mixer-layers. first accounts for the red cells in the diagram, that is, the direct dependency on previous layer's activations at position (line 7. thus finalizing by. After computing the activation at position (line B). 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 to be the same as before but within each the D dimensions that is:
Tu, ('r Σ Yi + Pt-i
for l' stsr'. The accounting works by considering (in-place) the influence of last positions of next U ones of b.
Algorithm 2 Flash Inference for LCSMs
Require: filter p E RMXLXD (lock[L.M first token af and sampler
1: Output: All activations a ...., M RR obtained by autoregressively sampling
2: Initialize 5 I Mx to zeros
to L do 4: U < maximum power of that divides # the side of the th gray tile to .M do 6: # account for the contribution of af é-1 to bf red cell
3: for
5: for
7:
if += af ④ p
8:
of block (bf)
9:
# account for the contribution of of Ui-- to 11411441-880 gray tile
10:
M(41540 += 7af-1
1, 1.+1i+40
11: end for
12: # generate next token based on the output of last layer at position
13: a sampler(a.))
14:
end for
As far performance goes, Algorithm practically and within each of the dimensions of embeddings). Proposition 2 covers its complexity analysis:
Argaritim calls Algorithm I M D times (within each layer
The overall time complexity of Algorithm 2 is O(MDL L) for the mixer-part plus LM block calls, to generate L L 2" tokens. Furthermore, for each 0 S q < P 1, there are M - calls to ranges of length 29, The overall amount of memory takes to store the (as is the case for the naive lazy approach) and the peak memory usage does
Proposition 2.
activations is not change the asymptotics.
7

---

Published as a conference paper at ICLR 2025
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 the inner loop and performed in parallel across all layers at once, just after the loop- this is because all their inputs and outputs are disjoint. However, the red lines need to be performed sequentially since of is a function of bf which a function of af for every E.
Note that this optimization can be applied the eager and lazy approaches well (provided that the red tiles are accounted for separately from the gray ones). However, the benefit of extra parallelization practically relevant only settings that are not memory-bandwide 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
bound this is the case for
3.3 MEMORY CONSIDERATIONS
As follows from Proposition T the overall size of the inputs to (and outputs of) r in the proposed algorithm is only per iteration as opposed to the naive implementations (such as lazy and eager approaches) that access an average of (2) positions. This directly translates to significant data movement improvements.
log I.) that is, we access on average activation values at log I. positions
simple optimization of static memory is to never store b, but rather use of to store of until hitting line this will work because from that point onwards bf never used again and af was not needed (or available) thus far. Hence, the static memory overhead is minimal: at least in the L.SSM case, one directly aggregates contributions to the same activation tensor, thus not using any extra storage (on top of the inherent M LD 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 ODM, but can be dropped to O(LD) at essentially no cost by dropping parallelization for large tiles, as further discussed in Appendix F
Lastly, if one is restricted to only using GPU and the whole tensor of M x L x D cannot fit it, it is possible to drop the requirement to only storing M x (L/2) x D activations by dropping parallelization in the largest tile (which was alluded above). This, together with more granular analysis of memory usage is further discussed in Appendices and Appendix F
3.4 GENERALIZATION
We propose a framework called Flash Inference to generalize our fast LCSM inference algorithm, covered in Appendix due to lack of space. Crucially, we identify the main properties that being exploited by Algorithm 2 in order to apply our tiling one needs the mixers to be contribution-base (that is, mixer(y); is some aggregation of contributions of yj to position i) and these contributions to only depend on data available at step j. The latter is the reason we cannot apply Flash Inference to transformers since (q. k2) depends on qi which depends on yi-
EXPERIMENTS
We performed two kinds of experiment: one synthetic and one on the Hyena architecture (Poli et al. 2023). The synthetic setup defines all blocks be MLPs with a hidden dimension of 2D and GELU activation; it also simply sets 4241 as a²d plus some noise to avoid dependency vocabulary size since that out of the scope of our framework and will be negligible scale. Note that this can be viewed sampler: function from logits at the last layer and previous position to the next token's 's embedding. In both settings the weights are initialized to random noise since it does not affect the runtime- this avoid unnecessarily training
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 E {1, 2, 4, 8), MM E 36), D E (256, 768} and generate tokens up the greatest power of 2 that fits the memory. All results are obtained
8

---

Published as a conference paper at ICLR 2025
by averaging over 4 runs following 2 runs of warm-up. We evaluate approach on the latest NVIDIA H100 and A100 GPUs. (Results for all the parameter sweeps in supplementary material).
For our basclines, (1) consider the cager and lazy approaches described in Section 3.1 depicted in 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 т, covered in Section 4.2 Our
best method is a Hybrid that dynamically chooses the optimal implementation of depending (B, D, M,MM.
4.1 INTEGRATING FLASH INFERENCE IN REAL WORLD SETTING (HYENA ARCHITECTURE)
Flash Inference significantly speeds-up the end-to-end inference by up to 7.8x and the convolution- based mixer component of the Hyena architecture by 110x compared to the baselines shown in Figures 2a and 2b Figure 2c 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 computed. For a given sequence length of L, we have L/2 positions that use tile size I, L/4 positions using tile size 2, and That is, 93.75% of tokens use tile size U S 8. Hence the spikes occur rarely.
ayer
Cornotee or putt butsh sicu 700000 10 layan, Leg!t
layers, Lego
127
Euger
Lasy
Euger Man aurulld
Lary Parsil
30'
AnD Teth Time-
Hob 10
59-
Enger
100
16-1
(Jay Non Horaie
164
107
54
1.kk
1.5
25a
Tarn sumber
Tooan Number
(a)
(b)
(c)
Figure 2: Real world Hyena experiments: (a) End-to-end inference time breakdown shows Hybird provides 4.8 speed-up over optimized baselines (b) Cumulative mixer time of Hybrid scales 90x better (c) Hybrid shows low variance in per-token response time except at the tokens positions where large tiles are computed.
4.2 T IMPLEMENTATIONS
Algorithm 2 assumes an implementation of primitive that accounts for the contribution of range of inputs to range of outputs of convolution. We considered 7 different implementations but, here, we only present results of the ones the Pareto Frontier that is, those optimal for at least some (B, N, D, U) setting (the others are covered in Appendix D.I). There are 4 such implementations, of two types:
(1) Sppthwise-separatee ID Convolution: We use two types of implementations (a) ConvID refers to the default Py PyTorch (Paszke et al. 2019) kernel, that requires explicit padding (b) Flash ConvID refers to fused kernel by FlashFFTConv (Fu et al. 2023b) these are asymptotically quadratic in tile size U.
(2) FFT based convolution: we again use two implementations (a) Torch native, shown as FFT, that requires computing the FFTs of inputs, followed by pointwise muitiplication and lastly performing the inverse FFT operation. (b) the second is provided by FlashFFTCom that does all the above in single fused kernel shown as FlashFFT- both these scale as OU U))
4.3 MIXER ISOLATION STUDY AND HYBRIDIZATION OF T IMPLEMENTATIONS
We evaluate the different convolution implementations for all settings of B, D. M, U and observe that each these four implementation lies on the pareto frontier curve of tile size vs latency as shown in Figure 3al Our Hybrid approach dynamically chooses the best implementation for given tile size U

---

Published as a conference paper at ICLR 2025
based on the isolated empirically impiriealllymmusurreed efficiency of each implementation shown in Figures Figure 3e shows the cumulative token time breakdown for the mixer and non-mixer components of the synthetic setup for all our implementations. We observe the non-mixer component staying largely the same across approaches this is made possible through the employment of CUDA Graphs further discussed in Appendix D.2
Me-38 D-228
butervscee DWCTA Lagie
Te' Broukthonn for buichaiee secr De lepe's, u42
Bith Care12
lapel
Oater
Rage Re
150
Core
21"
20
01
327
(a)
(b)
(c)
Figure 3: Mixer Isolation a Synthetic setting: (a) Different implementations are optimal for different tile sizes creating pareto optimal curve for Hybrid to choose, (b) Cumulative mixer inference of Hybrid achieves the best of all Implementations (c) End-to-ena cumulative token inference breakdown
4.4 IMPROVEMENTS JUSTIFICATION
This significant speed-up can be attributed to:
(1) The significantly lower (L log" L) FLOPs required for our tile computation approach compared to the (2...) FL.OPs required by Fager and Lazy counterparts. This is shown Figure 2b where methods based on our tiling outperform by large constant the quadratic methods in terms of time spent on the mixer components.
(2) Drastically reduced activation memory access from S(LL) for Eager and Lazy to out OLL log L) tiling-based methods. This can be shown through the performance of Flash ConylD (Figure 3b) which outperforms lazy and eager by margin although it also performs ?(L²) FLOPs it does so in a more memory- memory-friendly and kernel-optimizable way.
(3) The dynamic choice of best implementation for given tile U hybrid outperforming any method using a fixed implementation (Figure 3b
(4) Our engineering contributions: first, the DFT for the convolutional kernel is pre-computed ahead of time for all 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 processing is parallelized across layers 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).
CONCLUSION AND FURTHER WORK
We propose a framework for performing inference certain autoregressive sequence models. Among such models, LCSMs are noteworthy: there, our framework provides an ( log² L) inference algorithm which, when run empirically, yields up to 8x end-to-end 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 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 filters are made data deta-dependent, and, while previous works Arora et al. (203)); Karami & Ghodsi (2024) 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.
10

---

Published as a conference paper at ICLR 2025
6 ACKNOWLEDGMENTS
This work was partially enabled by use of the Kempner Al cluster. Sham Kakade acknowledges: this work has been made possible in part by 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.
REFERENCES
Simran Arora, Sabri Eyuboglu, Aman Timalsina, Isys Johnson, Michael Poli, James Zou, Atri Rudra, and Christopher Ré. Zoology: Measuring and improving recall in efficient language models. arXiv preprint arXiv:2312.04922 2023.
Dan Fu, Simran Arora, Jessica Grogan, Isys Johnson, Evan Sabri Eyuboglu, Armin Thomas, Benjamin Spector, Michael Poli, Atri Rudra, and Christopher Ré. Monarch mixer: A simple sub-quadratic gemm-based architecture. Advances in Neural Information Processing Systems, 36, 2024.
Daniel Y Fu, Tri Dao, Khaled K Saab, Armin W Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models. arXiv preprint arXiv:2212. 14052, 2022.
Daniel Y Fu, Elliot L. Epstein, Eric Nguyen, Armin W Thomas, Michael Zhang, Tri Dao, Atri Rudra, and Christopher Ré. Simple hardware-efficien long convolutions for sequence modeling. In International Conference on Machine Learning, pp. 10373-10391, PMLR, 2023a.
Daniel Y Fu, Hermann Kumbong, Eric Nguyen, and Christopher Ré. Flashfftconv: Efficient convolutions for long sequences with tensor cores. arXiv preprint arXiv:2311.05900 2023b.
Marco Gaido, Sara Papi, Matteo Negri, and Luisa Bentivogli. How do hyenas deal with human speech? speech recognition and translation with confhyena. arXiv preprint arXhv:2402. 2024.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752, 2023.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces, arXiv preprint arXiv:2111.00332 2021a.
Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Combining recurrent, convolutional, and continuous-time models with lincar state space layers. Advances in neural information processing systems, 34:572-585, 2021b.
Albert Gu, Karan Goel, Ankit Gupta, and Christopher Ré. On the parameterization and initialization of diagonal state space models. Advances in Neural Information Processing Systems, 35:35971- 35983, 2022.
Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces are effective as structured state spaces. Advances in Neural Information Processing Systems, 35:22982-22994, 2022.
Mahdi Karami and Ali Ghodsi. Orchid: Flexible and data-dependent convolution for sequence modeling. arXiv preprint rxXiv22402188008 2024.
Yuhong Li, Tianle Cai, Yi Zhang, Deming Chen, and Debadeepta Dey. What makes convolutional models great on long sequence modeling? arXiv preprint arXiv:22 10.09298, 2022.
Xuezhe Ma, Chunting Zhou, Xiang Kong, Junxian He, Liangke Gui, Graham Neubig, Jonathan May, and Luke Zettlemoyer. Mega: moving average equipped gated attention. arXiv preprint arXiv:2209. 100555....00 2022.
Xuezhe Ma. Xiaomeng Yang, Wenhan Xiong, Beidi Chen, Lili Yu, Hao Zhang, Jonathan May. Luke Zettlemoyer, Omer Levy, and Chunting Zhou. Megalodon: Efficient 1lm pretraining and inference with unlimited context length. arXiv preprint 2488002404.0880) 2024.
11

---

Published as a conference paper at ICLR 2025
Stefano Massaroli, Michael Poli, Dan Fu, Hermann Kumbong, Rom Parnichkun, David Romero, Aman Timalsina, Quinn McIntyre, Beidi Chen, Atri Rudra, et al. Laughing hyena distillery: Extracting compact recurrences from convolutions. Advances in Neural Information Processing Systems, 36, 2024.
Harsh Mehta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. Long range language modeling via gated state spaces. arXiv preprint arXiv:2206. / 13947, 2022.
Tom Le Paine, Pooya Khorrami, Shiyu Chang, Yang Zhang, Prajit Ramachandran, Mark A Hasegawa-Johnson, and Thomas S Huang. Fast wavenet generation algorithm. arXiv preprini arXiv:1611.09482, 2016.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models. arXiv preprint urXiv:2302. 10866, 2023.
David W Romero, Robert-Jan Bruintjes, Jakub M Tomezak, Erik J Bekkers, Mark Hoogendoorn, and Jan C van Gemert. Flexconv: Continuous kernel convolutions with differentiable kernel sizes. arXiv preprini 10.0v22110008059 2021a.
David W Romero, Anna Kuzina, Erik Bekkers, Jakub Tomezak, and Mark Hoogendoorn. Ckconv: Continuous kernel convolution for sequential data. arXiv preprint arXiv:2102.02611, 2021b.
Jiaxin Shi, Ke Alexander Wang, and Emily Fox. Sequence modeling with multiresolution convo- lutional memory. In International Conference on Machine Learning, pp. 31312-31327. PMLR, 2023.
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler, Efficient transformers: survey, ACM Computing Surveys, 55(6):1-28, 2022.
Aaron Van Den Oord, Sander Dieleman, Heiga Zen, Karen Simonyan, Oriol Vinyals, Alex Graves, Nal Kalchbrenner, Andrew Senior, Koray Kavukcuoglu, et al. Wavenet: A generative model for raw audio. arXiv preprint 160.0399, 12, 2016.
Joris van der Hoeven. international symposium on Symbolic and algebraic computation, pp. 17-20, 1997.
Lazy multiplication of formal power series. In Proceedings of the 1997
Joris Van der Hoeven. Relax, but don't be too lazy. Journal of Symbolic Computation, 34(6):479-542, 2002.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
A CONTEXTUAL DISCUSSION
In this section, extend our discussion the contextual relevancy of LCSMs and our method. We focus on the SISO (single-input-sinll--uuup such primitive is used across channels of the embedding dimensions form mixer.
setting unless otherwise specified and assume that
As discussed in Section 2.3 LCSMs directly compute the filter p to used to map an input sequence y E _ to an output sequence z E _R via 2, 2 Ei= 9: Pt-i. We assume here that E R' although it could be longer since no entry past 1 position matters. While is often underparameterized. one could, in principle, directly lear its entries (and, indeed, (Fu et al.) (023a) successfuly do that by using some extra regularization). However, even in underparameterizee settings, where p = f(0) for some O with dim(O) < L,, what filters p we end up dealing with is
12

---

Published as a conference paper at ICLR 2025
function of data and f: here, f implicitly bakes in certain bias. While depending bias (such as decaying pattern of p as suggested by eLi al. 2022)) ) one may assume to have certain exploitable structures, method is broadly applicable regardless of such structural properties. That makes Flash Inference be bias and data-independent. For completeness, we mention in Section (A.2 previooss works exploiting such structure that is, ways perform fast inference assuming certain properties of p hold.
Since, as mentioned in Section 2 there is direct connection between long convolutions and LTI SSMs, we cover this more extensively in Section A.
A.1 STATE SPACE MODELS
While SSMs and LCSMs can fundamentally represent the exact same class of models, they differ through their parameterization. Works such (Poli et al. 2023 Li et al. [022 Romero al. 2021b a Shi et al., 2023 Fu et al. 2023a) try to characterize the convolution implicitly. On the
other hand. LII SSMs (Gu et al. 202 2022 Ma al. 2022 2024 Gu al. 202 Gupta et al.,
2022, Mehta et al. 2022) define a recurrent system that is equivalent to a convolution and provide
an alternative perspective in doing We define discrete SISO SSMs below following the notation of (Gu et al. 2021a).
State space models (SSM) be thought of as linear RNNs. In its simplest form, SSM's mixer for a SISO operator mapping E R to 2 E Rh is defined via an extra sequence of hidden states w: R L'D' where D' is the SSM's dimension. For fixed input RR4, one defines:
U4 A(t)u-11+ B(t)ye
21 = C(t) ue
where A(I) D'DD' C(I) RD', While models like Mamba (Gu & Dao 2023) let A, B, C be functions of y and, thus, vary with linear-time invariant (LTT) models assume A, B and C/ to be fixed parameters, Thus, the equations rewrite to:
Ut Au By
2t = C. U1
Opening up the recurrence, we get that:
2t Σ (CT (CTA'=))-9) 1=0
which corresponds to a convolution with the filter p implicitly defined via P: = C AT If A is further assumed to be diagonal, we call the corresponding SSM diagonal. Gupta et al. (2022) showed that one does not lose practical model quality by assuming A to be diagonal. When one tried to perform inference on an LTI SSMs. one could use:
the recurrent view: keep around only the last state u4-1 and update it in-place. This is why the memory of performing inference with an SSM scales with D' rather than L. The time complexity of performing one step is O(D22) in the general case, but as mentioned, one can resort to diagonal SSMs which drop this complexity to O(D)).
convolutional view: precompute the equivalent convolution o and use Flash Inference.
A.2 COMPARISON TO STRUCTURAL FILTER EXPLOITATION
Following our discussion in Section A. L.TI SSMs become candidate of structural property that can be assumed on filter p R While if one to let /'' grow as much as T,, it is provably possible to represent any filter p, it is structural assumption to assume that a given filter p can be represented (or well-approximated) by an LTI SSM of a particular dimension D' (and this is what (Massaroli et al. 2024) do). However, if the filter has this structure then we use the recurrent view and pay a time cost of D' per token generation (assuming a diagonal SSM) while having our activation space drop by factor of L/D D'. Alternatively, we can store all the activations and use Flash Inference to get average token step for O(log" (T.) time-budget. Hence, at least asymptotically,
13

---

Published as a conference paper at ICLR 2025
if D' E Ollog(()) one would prefer the recurrent view (and thus using the D'-dimensional SSM equivalent to p). Past this point, our method is faster but still uses memory that scales up with L rather than D' so which one to choose is a function of our memory constraints.
Another noteworthy line of work follows dilated convolutional models and was popularized through (Shi et al. 2023). inspired in turn by (Van Den Oord et al. 2016). Paine et al. (2016) show how such models allow for linear time inference, while also needing to store all activations. Thus, granted such structure, Flash Inference is inferior and (Paine et al.) 2016 approach should be used.
Finally, it is worth emphasizing that Flash Inference can accommodate data-dependent filters, while online exploitation of structure is often prohibitively expensive (as in the case of distilling p into an LTI SSM following Massarolli et al., 2024h). Moreover, the class of data-dependent convolution- based models larger than that of LTI SSMs simply because of the linear-time invariant assumption: one cannot know the whole A, B. C when only half of is unlocked (similarly to how data-dependent
SSMs, such as (Gu & Dao (023). can represent more than LCSMs can). And while there are no
major data-dependent causal LCSMs yet, works such as (Arora et al. 2023 Karami & Ghodsi 2024) show the potential of such a research direction.
A.3 ON FLASH INFERENCE'S PREFILLING AND AUTOREGRESSIVE GENERATION STAGES
We have touched briefly in Section 2.3. how our method deals with prompts. In particular, if one is given a prompt of size P, FFT will be employed to fill in all the contributions of y1,P) to 2[1,))- The way this translates to the complete architecture is computing all of a [[.PP the contributions of these to the future L - P positions. One can do so at fraction of the cost of the forward pass, by just caching the appropriate activations and then zeroing out positions in the range [P + 1, L. before proceeding to the next layer (to not allow second order contributions to be performed). This comes a cost that scales with (L L).
[0.M] well as
Afterwards, we can drop the first P positions and apply our method starting with the current values of a to perform autoregressive generation. This is because the eager nature of our prefill is more aggressive than in the case of transformers since we already account for all the relevant contributions of the prompt to the future. This can be very useful when we have a long prompt, relatively few autoregressive steps to perform and the autoregressive generation takes place on a different machine, since we save up valuable and communication costs.
We also want to emphasize, that Flash Inference focuses on optimizing the autoregressive generation part. However, the prompt is large enough compared to the number of tokens generated, this may have smaller relative improvement since the overall time is dominated by the prefill.
B THE FLASH INFERENCE FRAMEWORK
We propose 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,
can get fast inference following the same kind of tiling.
Notice that in the case of LCSMs, never mattered how one computes the contribution of a range of inputs 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-speciffe but rather the way we tile the space of contributions while ensuring autoregressive generation can take place.
B.1 ARCHITECTURAN PROPERTIES
In order for our framework to apply, we need the mixers involved have certain properties:
P.1 Contribution The used mixers work by aggregating contributions of each input posi- tion to each subsequent output position. That is, for any mixer there exist an associative aggregation function agg X X, read function read X' IR and contribution
14

---

Published as a conference paper at ICLR 2025
function cont RR x NNNNN x N + X' such that:
mi)er(y), = read (agg((ont(() 1,),,,nt((4 2, conty i,
(6)
where, X' is a set of intermediate states and read a function to map those back to embeddings. Recall that agg being associative means that agg(z), T2, T.. agg(agg(T1,- .....,............... I.), agg(F=+1,..... .x,) for any S < т. For a sensible architecture,
the size of X and cost of agg should be of order D.
In the case of self-attention this translates to having read agg simulate the softmax, by letting X' = _D x R. agg = > and
cont i, = (Vy. e vi e QK" vi= - (r; (...)) (hh.))))
that is, the exponentially weighted value vector along with the exponential weight. Finally one can use read(v, w) v/w to implemnent the softmax normalization step. In the case of LCSMs, one can simply choose A' = RR read be the identity function, agg the sum again and cont (y, i,j) = Yi O Pj-i-
P.2 Query-independee
The contribution function cont j) does not depend on y[++1,4] Note that this is the case in LCSMs since y/ O Py-s only depends on V/- However, it is not the case for transformers, since cont(y. i.j) depends on qj which depends on yj.
B.2 SETTING AND DEFINITIONS
P.I holds and there exists algorithm A that for any given input sequence y and indices 1 I Er', ccomutees the contributions of U(I.r] to outputs at every position E [.'r r']:
Suppose
A(y, [rr 1' D'r agg(cont (y, (),, co, ( + L,p), ... coop.....oon((( cont r)))) Furthermore, for this choice of A, let x + such that for any S < S sr', evaluating Aly, ("', [","""")) r']) takes at most T(r - 1 + 1, r - 1) FLOPs.
we dropped the < l' condition and represent the procedure one needs to perform forward pass during training which refer as the static setting. In the case of LCSMs, A the FFT-based algorithm underlying Lemma]] with associated T T(L,, L2) = D(L, L2) log(L1 + L2), opposed to the naive implementation that would have Teaive L1, L2) = DL1 L2.
1' and
I., this algorithm would actually
B.3
MAIN RESULT
Our main result can be phrased follows:
generate L 2P tokens autoregressively by performing, per layer, - black- calls to A as well as more calls to cont, agg, read and block. Concretely, there are L/2 2P-1 calls to A of length each, L/4 2 2-2 calls of length each and so on up to call of length L/2. Hence, neglecting the cont, agg, read and block part, the overall time complexity per mixer layer is:
Theorem 2. Under the assumptions P.I and P.2
FLOPs- Σ
Furthermore, during every diaeenngnnration iteration, the calls to A across different layers can be performed in parallel there is data- -eta-dependyncy between them.
Here, we are being ignorant of the calls to cont, agg, read and block because they only scale linearly with however MLPs scale quadratically D and therefore these can be the limiting factor when D is large in comparison to L.
Proof. Figure shows why the tiling used in Algorithm 2 covers every pair of contributions ex- actly once and in the correct order (remember, we only assumed agg be associative, not nec- essarily commutative). The more general algorithm is illustrated in Algorithm B and follows the same shape its LCSM counterpart. The only difference is that 2 is not assumed to be
15

---

Published as a conference paper at ICLR 2025
R. necessarily, so one cannot store the intermediate accumulated states in the place they store the activations (though they could reuse some the space since af and by never need to coexist). We use be store the inner part of Equation agg(cont(a" ,,,,,,,, 1) cont(a
namely by incrementally computes ... cont(a 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 can be done parallelly across layers.
The only part remaining to be proved is the right counting of calls to A per layer: calls of length 29 happen whenever 29 divides but 2 does not that is, for E (29,3 3- 3,,,,,2,00 24,5""" 5.12,..... (2P 11)1)) 22). so 2 2- calls.
Algorithm 3 Generic Flash Inference
Require: Functions A, cont, agg, read, (1.ck(1.M]] first token of and sampler
Output: All activations a", ......, a Lx obtained by autoregressively sampling
2: Initialize 6 X with elements neutral to agg
3: for + L do
4: U f maximum power of that divides # the side of the 8th gray tile
5: for 8 + 1 to M do
6: # account for the contribution of aj af to bf red cell
7:
bf agg cant(af-1, i))
8:
block" break'(read(15) (b))
9:
end for
# account for the contribution of -i:-1+1.1 to 1141,4411 gray tile(s)
11:
parallelly across & +- M do: 1...0u A(a--, U 1 UJ)
12:
b 01+1:+0
# generate next token based on the output of last layer position
14:
U sampler(a,"
end for
C EXTENSION TO DATA- DEPENDENT FILTERS
The reason why data-dependent filters are harder to deal with is that cont(y, 1/4 Py-+ depends on both yE and on what Pj-v depends on. By only assuming causality, we can only access Pj-i after 2j-i-1 has been computed. This stops Algorithm 2 from working out-of-the-box since when is a power of 2, in order to account for the contribution of a1,1 to af 1 one will need access to of 1122111 which is guaranteed to be accessible (only pf..q be assumed be known this stage).
The modified algorithm is shown in Algorithm 4 and precisely follows the tiling of van der Hoeven (1997) and, thus, its correctness transfers. Note that rather than using the 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 rectangle, although the other way around is possible (and is what we did through Lemma IT. This implementation performs convolutions between two sequences length U each and uses the whole output and thus requires order 20 FFT. Note that this happens twice for given i when i + 1 is not power of 2 (almost always). Algorithm 2 on the other hand, only performs a convolution between a length-U sequence and length-2U sequence. However, noted in Appendix [D] we can get away without padding and thus performing only one order 20 FFT. Thus, our proposed tiling improves FLOPs by a factor of 2, but it requires the kernel to be data-independent.
D IMPLEMENTATION IMPROVEMENTS
In Algorithm 2. the calls to all have shape of T U 1, p. 1, where U is power of 2 that divides i. Following the proof of Lemma we note that to implement т. we need to perform the convolution of segment of length U of y and prefix of length 20 of p. However,
16

---

Published as a conference paper at ICLR 2025
Algorithm 4 Flash Inference for LCSMs with data-dependent filters
Require: first token dn db. Pi.. block (M.) and sampler
Output: All activations a ,...,, a
4x obtained by autoregressively sampling
2: Initialize a with zeros outside 0g
3: for + L - do
4: U < maximum power of that divides (i 1)
5: for 8 < to M do
6:
Compute pf as causal function of data [
f-1
per model specification
7: # account for the newly available contributions
8:
of block" (af aaf- pá a6-1 @ p$)
9:
# account for some cager contributions
if + 1 U then
11:
# downgrade the power of 2 by half if + 1 is already power of 2 U +U/2
13:
([U,AUU2] += CONV 11: P(u,20-1)
14:
else
15:
of (1+1,42U-1) CONV 6-1 1U.20-11 (O P ((411)))
16:
114 CONV (M1,2U- up (--U11))
end if
18:
end for
# generate next token based on the output of last layer at position
20: 0 sampler(a."
end for
following the convolution, of the 3U 1 outputs, indexed 3U 2]. we only interested in the middle of them, namely indices [U, 2U 1]. The canonical way of performing convolutions via FFT involves padding by enough Os and using order large enough store the whole output which rounded up to the closest power of 2 would mean an order AU FFT. However, using 2U FFT call, which will perform 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 U - 2]. we do not intersect [U, 2U - 1). Furthermore, there are only log L different lengths of prefixes of p involved in these convolution so one could, in fact, cache the DFTs for these as precomputation step, thus dropping the number of DFT's per convolution from
speeding up by further x1.5 factor.
D.1 MORE IMPLEMENTATION DETAILS
In Section 4.2 it has been mentioned that we considered 7 implementations of Besides the four mentioned (Torch's and FlashFFT's Conv ID and FFT) we also did 3 more naive versions:
1. Torch FFT (no precomputed FFT for the convolution kernel): This implementation is identical to Torch-FFT except that we do not pre-compute the DFT of the convolution kernel beforehand.
2. Vectorized Conv1D: It implements 1D convolution operation in vectorized manner using tensors and leverages efficient operations to process multiple slices in parallel. It takes advantage of the fact that convolutions depth-wise separable.
3. Element-wise Conv ID: This implementation does element-wise dot product of the convolution kernel with the input, while iteratively sliding the kernel over the input.
D.2 CPU OVERHEAD: THE SILENT PERFORMANCE KILLER
One primary improvement introduce is the of CUDA Graphs to mitigate the CPU overhead observed in the non-mixer components during inference.
In both synthetic settings and the Hyena model, these non-mixer components primarily deal with only current token being generated. As a result, the GPU compute time is often shorter than the kernel dispatch time, leading CPU overhead. Typically, this overhead remains hidden when the
17

---

Published as a conference paper at ICLR 2025
mixer component requires significant GPU compute time. However, Flash Inference dramatically reduces mixer times (up to 120x), exposing the CPU overhead and causing the observed decrease in end-to-end speed-ups. A similar effect is observed for the Lazy approach compared to Eager, as Lazy's 's faster execution also exposes the overhead.
To address this, we employ CUDA Graphs, which record all kernel dispatches for generating single token into a graph. For subsequent token generation, the same graph is replayed, allowing all kernels for the non-mixer components be dispatched simultaneously. This approach significantly reduces CPU overhead and ensures consistent overhead times for both Flash Inference and its baselines.
With these improvements, Flash Inference demonstrates much better end-to-end relative gains, the mixer now constitutes larger portion of the total runtime. We observe speed-ups of up to 120x in the mixer and up to 8x in the end-to-end Hyena runs (as opposed to 1.6x when not using CUDA Graphs). All the results are shown in Appendix !
E STORING ONLY HALF OF THE ACTIVATIONS
If the whole activation tensor (M L x D) cannot fit the memory of one GPU, we save factor of 2 along the L-axis. To do this, observe that after (L/2)th iteration completes we never need to look back any activation position i < 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 A work as an in-place operator requires storing M LD intermediate values, one can choose not apply 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 A is read, we place 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 A can be LD or even
As we get the output of A(a", 1, L/2, L/2 11, L),, which has a shape of (L/2) x D, we overwrite a by and then move on to the next layer. We can do this because the contribution of aς to at has already been properly accounted for so the first half of a" truly becomes irrelevant for next layers and iterations to come. Hence, a" contains the intended second half of a We then proceed to do the same thing for a! and so on. At the end, a0 will contain the second halves of a!!...M), so we can shift these by 1, emptying a" and now finally having a represent the second halves of activations the first halves have been discarded.
Although per layer, one will have peak memory containing both a², the output, and whatever other necessary temporary values needed to perform A, this does not scale with M because we essentially reuse the extra space needed by A across different layers. We seemingly pay the cost of not performing the calls to in parallel across layers, but if storing the M x L x tensor of activations an issue, then one would very likely have to make this sort of compromise to be able to run A in the first place.
F DISCUSSION ON MEMORY
Besides peak memory usage, it is important to understand how memory boundedness affects end-to- end performance. In general, GPUs have types of memory bound:
Memory Hetency--ouudee most time is spent on memory operations, but parallelism is insufficient to saturate GPU memory bandwidth. In this setting, throughput can be optimized by increasing parallelism saturate memory bandwidth.
Memory Bandwidth-Bound parallelism. In this case, throughput optimization is achieved by increasing arithmetic intensity (compute-to-le ratio), usually by changing the algorithm or fusing operations.
GPU memory bandwidth is already saturated with available
Putting this into perspective for our use cases, and looking at how across-layer parallelization affects performance, we have that:
Lazy computes token contributions using pointwise multiplication kernels, accessing almost all sequence activations for large values, making it memory bandwidth- bound beyond
18

---

Published as a conference paper at ICLR 2025
certain (early) iteration index. Hence across- layer parallelization helps most in the beginning. Lazy acts exactly as a dual (it is memory bandwidth- bouud until close to the end).
Our approach computes token contributions in a tiled fashion, with most tiles being small (87.5% have size < 4), making the workload memory-latenes bound. Parallelizing across layers can then further saturate memory bandwidth. For large tile sizes we end up using FFT implementations decrease arithmetic intensity since we get to a compute-boune regime.
In Section 3.3 we discussed the peak memory usage being given by the gray calls to т. Since each call to T performs D FFTs of length each (when dealing with the biggest tile of side L/2),, the naive implementation would use up MD 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(L)). or even just (L) if one treat sequentially the different dimensions. Dropping parallelization could theoretically yield a slowdown, but that is only when we are not memory bandwidth- bound which the case one has to worry about allocating more than (I.D) in that case, the FFT calls are anyway not happening in parallel. Hence, one can easily drop the peak memory overhead to OLL)) or even (L) without incurring an actual time cost.
G FRAMEWORK ADAPTATIONS
We covered the clean case of LCSMs and of other architectures that fit Theorem However, all these assume complete causality of the architecture as well as all the layers fitting our framework, This need not be the case and Flash Inference can easily be adapted to deal with both.
As far as causality is concerned, one needs to assume for autoregressive inference to be dell-defined. Remember, we defined causality by having af (xx), be function of only Tp1 or, equivalently, (ixer'(y); is a function of VI... If this was not the case, then we could not compute I4+1 by only knowing z4. However, much like the encoder-decoder architecture, we can allow for non- causal prompts. That is, if at any point an external user gives us a sequence of tokens (say an image), we could fill in the activations at those position in the standard forward-pass way, and then start our procedure from there as if we were given a fixed a prompt. Granted a prompt and the can proceed with
activations at its respective positions, Section [23.I] following (Massaroli et al. 2024 Lemma 2. 2), and essentially accounting in ad advance for the contribution of all this prompt all subsequent positions.
inference method as described in
Another useful adaption for most alternatives to attention is that of using hybrids: architectures where different layers use different mixers. Our method can be applied seamlessly to the subset of layers that fit the framework. All the algorithm stays unchanged but computing a (xx); for the layers é which do not fit our framework needs to be done in the way specific to mixer". Importantly, even the across-layer parallelization can be maintained since the gray tiles only affect subsequent positions at the layers that fit the framework, hence not interfering with the rest. At any point in our algorithm, we have computed every piece of information that the lazy approach has, as well as some extra partial future contribution. Therefore, everything that worked assuming the lazy approach will work with our method as well this includes other layers attending to more than the previous layer.
H PROOFS
Lemma. Let < < L represent ranges of lengths L1 r and of y and z, respectively. There exists an FFT-based algorithm running in O(L1 + L2) space and (( + I2) log(T) + T22)) time complexity that, given access to U[..) computes all the elements of T(y 1rr,,,. that is, all the aggregated contributions of Vit.)) to all of z'1
Proof. Consider performing convolution between y[( and Pp-- It holds that for each j E r'] and i E [l, i is in the range of p. On the other hand, truncating the output of the convolution appropriately, one can keep only the corresponding outputs for each j E [['r'], The overall size of the FFT is then given by the sum of the lengths, which is at most - 1 + 1 (r'
19

---

Published as a conference paper at ICLR 2025
(((-))+)((++++ (r'-+++)) Since an FFT
of order runs in Oog L) and takes up O(L) memory, the conclusion follows.
I EXTENDED EXPERIMENTAL RESULTS AND ANALYSIS
1.1 BAR PLOTS
For completeness, provide the breakdown for both Hyena(Figure 1 and synthetic(Figure 5 settings. As mentioned in the main body, we see consistency across the non-convolution part and very significant improvement the convolution part.
1.2 SUMMARY
We provide the end-to-end relative speedups for our methods with regards to Lazy (the parallel version) for all the 21 runs that we performed, for Hyena as well as the synthetic setup, measuring both complete wall-clock time, well mixer time.
Notice that some the synthetic runs, torch-based FFT implementation is faster than Hybrid attribute that to noise since the absolute runtime difference is below 50 ms and it translates less than 2% drop in performance. Hybrid however sometimes exceeds FFT's performance by close to 20% (for example in the very first experiment).
Worth noticing also the consistent 10% speedup that the parallelization of Lazy yielded do note that the reported speedups relative to the parallel version of Lazy, but they relative to the non-parallel version. Furthermore, Conv ID yields a consistent > 5-fold improvement while still having quadratic complexity this shows the potential of our framework as a grouping mechanism even when no asymptotically faster algorithms to deal with the tiles exist: we simply get better arithmetic intensity.
bigger
'The Conv ID kernel crashes for configurations B = N,N = [18, 368,6], L - - This is because it
requires 32-hit indexing which is only supported in cu/NN version 9.3 and above. Our experiments use cuDNN 19.
20

---

Published as a conference paper at ICLR 2025
Table 1: Mixer times improvements on end-to-end Hyena relative to the lazy approach (higher is better)
B
N
D
L
Hybrid
Eager
Eager_NP
Lazy_NP
18
256
2
95.96 x
0.54
0.52
0.85
768
2
100.38 x
0.54 x
0,53
0.89 x
864
2
110.74 x
0.53 x
0.53 x
0.91 x
36
256
2
79.71 X
0.54 x
0.50
0.76 x
768
26
89.36 x
0.54
0.51
0.82
864
2
94.11 x
0.54 x
0.52
0.84 x
2
18
256
2
75.22 x
0.54 x
0.50
0.85 x
768
26
84.88 x
0.54
0.51
0.87 x
864
2
93.35 x
0.54
0.51
0.86
36
256
2
64.60 x
0.54 x
0.49
0.77 x
768
2
74.80 x
0.55 x
0.50 X
0.80 x
864
2
81.20 x
0.55 x
0.49
0.77 X
4
18
256
26
66.82 x
0.54
0.53
0.88
768
2
78.79 x
0.55 x
0.53
0.87 x
864
2
84.77 x
0.55 x
0.52 X
0.88 x
36
256
26
74.92 x
0.54 X
0.49
0.79
768
2
88.22 x
0.56
0.51
0.84
864
2
92.96 x
0.56 x
0.49
0.81 x
8
18
256
2
81.37 x
0.55 x
0.53 x
0.90 x
768
2
96.14 X
0.56 X
0.53
0.90 X
864
215
104.81
0.56
0.53
0.91
21

---

Published as a conference paper at ICLR 2025
615
355
Lruif
227
300
340
DAy
Saper
(045) 1kys
Iaovite 2. 0H000 50 eans L-25
BX0X Di SeGipen Ligit
400-
L-g
Ligit
201
3:50
240
DAY
5009 AP 100
(ium 1967
100
2.0 3.0-220022000 epnt 1-21
Bech 2-256. 25 synn 1-21*
bwch L. 0-700 eenn L-21*
Derch 206 0-756 35 lyyess L-g1>
100
200
100000
60
112
170
125
57
NarkCam Rep
Rege
Ryand (1mn
1ap
byane Duril Repe
Figure 4: Time breakdown for end-to-end Hyena experiments.
22

---

Published as a conference paper at ICLR 2025
Baider micoon E D-38 Wllenogg
Figure 5: Time breakdown for end-to-end synthetic experiments.
23

---

Published as a conference paper at ICLR 2025
Table 2: Cummulative times improvements on end-to-end Hyena relative to the lazy approach (higher better)
B
N
D
L
Hybrid
Eager
Eager_NP
Lazy_Ny -PP
18
256
2
5.30
0.59
0.57 x
0.88
768
2
5.60 x
0.58 x
0.57 x
0.89 x
864
2
5.71 x
0.59 x
0.59 x
0.93 x
36
256
2
3.11 x
0.66 x
0.61 x
0.83 x
768
1
3.60
0.64
0.62
0.90
864
2
3.50 x
0.60 x
0.59 x
0.85 x
2
18
256
2'7
4.30 x
0.61 x
0.56 x
0.88 x
768
26
4.81 x
0.60 x
0.57 x
0.91 x
864
26
5.09
0.59
0.56
0.89
256
2
2.63 x
0.65 x
0.60 x
0.80 x
768
2
3.01 X
0.64 x
0.60 x
0.85 x
864
2
3.16 X
0.62 x
0.58 x
0.80 x
4
18
256
26
3.88
0.61
0.59
0.90
768
2
4.53 x
0.61 x
0.59 x
0.91 x
864
215
4.85 x
0.60 x
0.57 x
0.89 x
36
256
2
3.76 X
0.61 X
0.55 x
0.81 x
768
2
4.46
0.61
0.56 x
0.85
864
2
4.86 x
0.61 x
0.54 x
0.82 x
8
18
256
216
6.95 x
0.58 x
0.57 x
0.92 x
768
1
7.14 X
0.59 X
0.56 x
0.91 x
864
25
7.83
0.59
0.56 x
0.91
24

---

Published as a conference paper at ICLR 2025
Table Mixer times improvements on end-to-end synthetic relative to the lazy approach (higher is better)
B
N
D
L
Hybrid
CID
flashCID
flashFFT
FFT
Eager
Eager NP
Lazy NN
18
256
2
97.70
5.94 x
2.40
51.16
82.27 X
0.54 X
0.52
0.85 x
768
217
112.62 x
2.42
51.18 x
111.57 x
0.54 x
0.53 x
0.90 x
864
2
124.30 X
2.42
56.14 x
119.41 x
0.53 x
0.53 x
0.91 x
36
256
2
82.69 X
5.94 x
2.41
39.73 x
78.58 x
0.54 x
0.50 x
0.76 x
768
2
100.53
2.41
35.36
102.32 x
0.54
0.51
0.82
864
2
112.33 x
2.41
36.58
112.81 X
0.54 x
0.52 x
0.84 x
2
18
256
2
84.84 X
5,94 x
2.41 x
52.28 x
81.43 X
0.54 x
0.50 x
0.85 x
768
26
84.84
5.81 x
2.40
45.94
88.36 X
0.54 x
0.51
0.87 x
864
26
97.23
5,85 x
2.40
49.77
94.40
0.54 x
0.51
0.87 x
36
256
26
69.28 x
5.74 x
2.40 x
39.57 x
69.03 x
0.54 x
0.49 x
0.77 x
768
2
78.79 X
5,54 x
2.43 x
31.91 X
76.47 X
0.55 x
0.51 x
0.80 x
864
25
85,89 X
5.56 x
2.43
33,80 X
87.56 X
0.55 x
0.49
0.77 x
18
256
26
75.13
5.97 x
2.49
47.70
74.52
0.54
0.53
0.88 x
768
2
80.48 x
5.75 x
2.52 x
42.66 x
82.05 x
0.55 x
0.53 x
0.87 x
864
2
93,40 x
5.79 x
2.54 x
45.22 x
88.84 x
0.55 x
0.52 x
0.88 x
36
256
26
79.85 X
5.78 x
2.40
50.57 x
80.60 X
0.54 X
0.49
0.79 x
768
2
92.30
6.47 x
2.45
46.37
88.90
0.56
0.51
0.84 x
864
2
96.29 X
6.49 x
2,45
47.93 x
94.17 x
0.56 x
0.50 x
0.81 x
18
256
216
95.32 x
6.10 x
2.51 x
56.11 x
87.85 x
0.55 x
0.53 x
0.90 x
768
25
102.49 X
5.89 X
2.55
54.83 X
94.08 X
0.56 x
0.53 X
0.90 X
864
215
107.36
5.93
2.55
45.75
98.61
0.56
0.53
0.91 x
25

---

Published as a conference paper at ICLR 2025
Table 4: Cummulative times improvements on end-to-end synthetic relative to the lazy approach
(higher is better)
B
N
D
L
Hybrid
CID
flashCID
flashFFT
FFT
Eager
Eager
Lazy_NP NP
18
256
2
6.70
3,55
2.04
6.34 x
6.64
0.59
0.57
0.87
768
217
8.36 x
2.09 x
7.08 x
8.05 x
0.57 x
0,56
0.91 x
864
2
7.76 x
2.09 x
6.66 x
7.42 x
0.57 x
0.56
0.91 x
36
256
217
4.32 x
2.91
1.89 x
3.88 x
4.48 x
0.60
0.57
0.78 x
768
2
5.17
2.00
4.56
5.50
0.61
0.58
0.85
864
2
5.01 x
1.91 x
4.30 x
4.78 x
0.60
0.58
0.87
2
18
256
2
6.02 x
3.30 x
2.01 x
5.70 x
6.01 x
0.59 x
0.55
0.88 x
768
26
6.06
3.34
1.97 x
5.75 x
5.90 x
0.58
0.55
0.89
864
26
6.38
3.50
2.02 x
6.03
6.37
0.57
0.54
0.88
36
256
2
3.84 x
2.59
1.84 x
3.67 x
3.94 x
0.63 x
0.57
0.82 x
768
2
4.25 x
2.74 x
1.84 x
3.69 x
3.89 x
0.62 X
0.57
0.83 x
864
25
4.19
2.71
1.83 x
3.88 X
4.51 X
0.62
0.56
0.81 x
18
256
26
5.83
3.34
2.07
5,46
5.77
0.59
0.57
0.89
768
2
6.25 x
3.45
2.08 x
5.91 x
6.29 x
0.60 x
0.57
0.89 x
864
2
7.15 x
3,45 x
2.08 x
6.24 x
6.64 x
0.59 x
0.56 x
0.90 x
36
256
26
5.98
3.33
2.02 x
5,49 X
6.38 X
0.59
0.53
0.82
768
2
6.62
3.64
2.06
6.75
6.59
0.59
0.55
0.85
864
2
6.99 x
3.74
2.06 x
6.68 x
7.13 x
0.59 x
0.53
0.84 x
18
256
216
9.88 x
4.19 x
2.22 x
9.14 x
10.36 x
0.57 x
0.56
0.90 x
768
25
11.55
4.30
2.29 X
9.85 X
10.55
0.58 X
0.55
0.91 X
864
215
11.37
4.25
2.28 x
10.56
11.20
0.57
0.55
0.91
26


Reviews:
[{"review_id": "BtFXFBsK6X", "reviewer": "Reviewer_zv8x", "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.\n\nThe 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.", "questions": "Please see my question above.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "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.\n\n* 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).", "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.\n\nLong 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].\n\nCrucially, 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.\n\nGiven 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.\n\nReferences:\n\n[1] Van Den Oord, Aaron, et al. \"Wavenet: A generative model for raw audio.\" arXiv preprint arXiv:1609.03499 12 (2016).\n\n[2] Shi, Jiaxin, Ke Alexander Wang, and Emily Fox. \"Sequence modeling with multiresolution convolutional memory.\" International Conference on Machine Learning. PMLR, 2023.\n\n[3] Paine, Tom Le, et al. \"Fast wavenet generation algorithm.\" arXiv preprint arXiv:1611.09482 (2016).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5I9yN2WhX5", "reviewer": "Reviewer_iRJb", "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.", "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.\n\nHowever, 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\n- 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)\n- Showing that there exist data-dependent extensions that outperform data-independent convolutions in modeling performance", "limitations": "", "rating": 5, "confidence": 4, "soundness": 4, "presentation": 2, "contribution": 2, "strengths": "The paper's main strength is that it is technically sound and novel, and addresses the problem it aims to solve.\n- The technical writing is clear and is helped by the inclusion of helpful graphics and rigorous algorithm boxes.\n- Many considerations and variants of the core algorithm are proposed.\n- An actual implementation is provided and all variants and baselines are benchmarked empirically.\n\nThere 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.", "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.\n\n- 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).\n- 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.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "D3QbTXxpZy", "reviewer": "Reviewer_xHkL", "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.\nThe resulting algorithm results in important speeds-up both for the position-mixing components and end-do-end inference.", "questions": "### Additional questions and observations\n\n- 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.\n\n- Line 233. There’s a typo here that changes the whole meaning of the sentence. Please fix.\n\n ### Conclusion\n\nWhile 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. \n\nDo 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.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The method is novel and offers interesting improvements in the inference speed of LCSMs.\n\n- The paper offers interesting perspectives that could be used for the design of more efficient (causal, input-dependent) LCSMs in the future.", "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.\n\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5zOp1mNIbf", "reviewer": "Reviewer_qm9L", "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.", "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?\n\n2. 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?\n\n3. 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?\n\n4. 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.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 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.\n\n2. 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.\n\n3. 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.", "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.\n\n2. 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.\n\n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "M08Em5PBC3", "reviewer": "Reviewer_pkhg", "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.", "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?\n2. 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.\n2. 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? \n3. 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.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. Technical Innovation and Soundness:\n- The paper presents a mathematically rigorous approach to improving inference efficiency, with clear proofs and careful analysis of complexity bounds\n- The implementation details are thoroughly considered, including memory optimization and parallelization strategies\n2. Practical Impact:\n- The achieved speedups (1.6× end-to-end, 50× for position-mixing) are significant and well-documented\n- The method is exact rather than approximate, maintaining model fidelity while improving performance\n- The framework extends beyond just LCSMs, providing valuable insights for future architecture design\n3. Experimental Validation:\n- Comprehensive empirical evaluation across different hyper parameters parameters ($B$, $M$, $D$)\n- Careful ablation studies of different $\\tau$ implementations\n- Clear breakdown of performance improvements and their sources", "weaknesses": "1. Technical Presentation:\n- 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. \n- the use of \"$\\mapsto$\" notation on line 159 is imprecise; it should be $y\\mapsto z$.\n- Algorithm 1 should reference equation (3) for the definition of $\\tau$ for better clarity.\n\n2. Contextual Discussion:\n- 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.\n- 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.\n- 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.\n\n3. Technical Specifics:\n- The asymptotic complexity expressions should consistently specify the base of the logarithm ($log_2$)\n- The parallelization discussion could benefit from more detailed analysis of the memory bandwidth implications\n- 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?)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "sUSu2qHlxP", "reviewer": "Reviewer_byxw", "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.", "questions": "1. Please kindly provide more details to the questions raised in the \"Weaknesses\" section. \n2. 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? \n3. 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?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper is well-written and easy to follow. \n2. The complexity analysis regarding both the computation and memory is detailed. \n3. The experiments are also provided to numerically verify the effectiveness of the proposed method.", "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. \n2. 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)? \n3. The numerical experiments part can be strengthened by e.g. adding tests for multiple configurations of hyper-parameters.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)