FlashFFTConv: Efficient Convolutions for Long Sequences with Tensor Cores
Authors: Daniel Y Fu
Hermann Kumbong
Eric Nguyen
Christopher Re
Venue: ICLR 2024 poster (2024)
Keywords: convolutions
GPUs
hardware-efficient algorithms
long context
fast fourier transform
I/O awareness
Primary area: infrastructure, software libraries, hardware, etc.
TL;DR: We propose FlashFFTConv, a new system that optimizes the FFT convolution algorithm to speed up long convolutions and enable long-sequence applications.

Abstract: Convolution models with long filters have demonstrated state-of-the-art reasoning abilities in many long-sequence tasks but lag behind the most optimized Transformers in wall-clock time.
A major bottleneck is the Fast Fourier Transform (FFT)---which allows long convolutions to run in $O(N\log N)$ time in sequence length $N$ but has poor hardware utilization.
In this paper, we study how to optimize the FFT convolution.
We find two key bottlenecks: the FFT does not effectively use specialized matrix multiply units, and it incurs expensive I/O between layers of the memory hierarchy.
In response, we propose FlashFFTConv.
FlashFFTConv uses a matrix decomposition that computes the FFT using matrix multiply units and enables kernel fusion for long sequences, reducing I/O.
We also present two sparse convolution algorithms---1) partial convolutions and 2) frequency-sparse convolutions---which can be implemented simply by skipping blocks in the matrix decomposition, enabling further opportunities for memory and compute savings.
FlashFFTConv speeds up exact FFT convolutions by up to 8.7$\times$ over PyTorch and achieves up to 4.4$\times$ speedup end-to-end.
Given the same compute budget, FlashFFTConv allows Hyena-GPT-s to achieve 2.3 points better perplexity and M2-BERT-base to achieve 3.3 points higher GLUE score---matching models with twice the parameter count.
FlashFFTConv also achieves 96.1% accuracy on Path-512, a high-resolution vision task where no model had previously achieved better than 50%.
Furthermore, partial convolutions enable longer-sequence models---yielding the first DNA model that can process the longest human genes (2.3M base pairs)---and frequency-sparse convolutions speed up pretrained models while maintaining or improving model quality.

Published as a conference paper at ICLR 2024
FLASHFFTCONV: EFFICIENT CONVOLUTIONS FOR LONG SEQUENCES WITH TENSOR CORES
Daniel Y. Fu', Hermann Kumbong", Eric Nguyen, Christopher Ré "Equal contribution. Stanford University.
fu,,fu, kumboh, etnguyen, chrismzer @stanford.
ABSTRACT
Convolution models with long filters have demonstrated state-of-the-en reasoning abilities in many long-sequence tasks but lag behind the most optimized Trans- formers in wall-clock time. A major bottleneck the Fast Fourier Transform allows long convolutions to run in O N) time in sequence length N but has poor hardware utilization. In this paper, we study how to optimize the FFT convolution. We find two key bottlenecks: the FFT does not effectively use specialized matrix multiply units, and incurs ex expensive I/O between layers of the memory hierarchy. In response, we propose FLASHFFTCONY. FLASHFFTCON uses a matrix decomposition that computes the FFT using matrix multiply units and enables kernel fusion for long sequences, reducing I/O. We also present partial convolutions and 2) frequency-span
sparse convolution algorithms- convolutions-which can be implemented simply by skipping blocks in the matrix decomposition, enabling further opportunities for memory and compute savings. FLASHFFTCONN speeds up exact FFT convolutions by up to 7.93x over Py Torch and achieves up to 4.4x speedup end-to-end. Given the same compute budget, FLASHFFTCONY allows Hyena-GPT GPT-s achieve 2.3 points better perplexity on the PILE and M2-BERT-bass to achieve 3.3 points higher GLUE score- matching models with twice the parameter count. FLASHFFTCONY also achieves 96. %% accuracy Path-512, high-resolution vision task where model had previously achieved better than 50%. Furthermore, partial convolutions enable longer-sequence models yielding the first DNA model that can process the longest human genes (2.3M base pairs)-an and Siqquency-sparss convolutions speed up pretrained models while maintaining or improving model quality.
-oodels-yielding
INTRODUCTION
A key challenge in machine learning is to efficiently reason over long sequences. Recently, convolutions have emerged as key primitive for sequence modeling, underpinning state-of-the-art performance in language modeling (Ma et al., 2022; Poli et al., 2023; Wang et al., 2022a; Fu et al., 2023a), time-series analysis (Tang al., 2022: Zhang et al., 2022b; Gu et al., 2021; Fathullah et al., 2023), computer vision (Wang et al., 2023; Liu et al., 2022; Nguyen et al., 2022), DNA modeling (Nguyen et al., 2023), and more (Li al., 2022; Islam al., 2023; Kim & Park, 2023; David et al., 2022: Miyazaki et al., 2023; Mehari & Strodthoff, 2023). Despite these strong quality results-and other benefits ranging from better scaling in sequence length (Gu al., 2021) to greater stability (Tay et al., 2021 la; Bietti & Mairal, (007))-covoluutoone sequence models still lag behind Transformers in wall-clock time.
major reason is poor hardware support. Unlike classical convolutions used vision applications, which often have short filters (e.g., x or 7 (Krizhevsky et al., 2012: He et al., 2016)), convolutions for sequence modeling often use filters as long as the input sequence (Romero et al., 2021b; Li et al., 2022). Such long filters necessitate the use of the FFT convolution algorithm. which computes the convolution between an input u and convolution kernel & via a conversion to frequency space:
(w*k) k
Σ
uliikblji - 2 u -
(1)
where F is the FFT, which can be computed in (Nlog)) time in sequence length N, and @ is elementwise multiplication. Despite its asymptotic efficiency, the FFT convolution algorithm has

---

Published as a conference paper at ICLR 2024
Registers
T9/
FFT
FNT
A2
24
Unpad
Frequency Time
64 KB
Pointwise
QAA.
Time (ms)
Pointwise
20 MB
SRAM
TB/s
Fused
HBM
IFFT
40
Pad
GPU Memory Hierarchy
Order-p Monarch Decomposition FFT Convolution on Tensor Cores
PyTorch FlashFFTConv
Sparse Convolution Filters
Speedup Through Fusion
Figure I: Left: GPU memory hierarchy. Middle left: Order-p Monarch decomposition of FFT convo-
lution, with 2. Middle right: Kernel fusion for end-to-end speedup. Right: FLASHFFTCONY introduces analogues of sparsity for convolutions.
poor wall-clock time on modern accelerators. In contrast, systems advances have pushed Transformers to the limits of modern accelerators achieving more than 72% FLOP utilization end-to-end with FlashAttention-v2 (Dao et al., 2022b; Dao, 2023).
In this paper, we study how optimize the FFT convolution algorithm on modern accelerators, to enable longer-context abilities. Just as systems advances such as FlashAttention yielded improvements in modeling quality (Ahdritz et al., 2022; Li et al., 2023) and the development of new attention algorithms (Paliotta et al., 2023; Athiwaratkun et al., 2023; Kwon et al., 2023; Liu & Abbeel, 2023), we hope that understanding how to optimize the FFT convolution can also inspire algorithmic innovation, thus improving the quality of convolutional sequence models.
For short sequences, the FFT convolution relatively easy to optimize. Kernel filters are often shared across many batches, which allows prepconpu the FFT of the filter k Fk and re-using it in batch: (u* = (Fu@k) Thus the FFT convolution is pleasantly parallel across batches and filters, and intermediate outputs the convolution can be cached in SRAM registers via kernel fusion.
However, as sequence length increases, we find that two key bottlenecks emerge. First, FFT convolutions do effectively use the specialized matrix-matri multiply units available on modern accelerators -e.g. the H100 can use tensor cores to compute matrix-matrix matrix multiply at 1.0 PetaFLOP/s compared 67 TeraFLOP/s for general arithmetic. Second, sequences become too large to fit in SRAM, and kernel fusion fails, resulting in expensive I/O costs (Figure 1 middle right). These I/O costs can be exacerbated by padding operations for causality, and conversions from real-valued inputs/outputs complex-valued FFT intermediates.
In response, we propose FLASHFFTCONY, a new system that optimizes the FFT convolution for long sequences using a Monarch decomposition of the FFT. An order-p Monarch decomposition rewrites the FFT as a series of p matrix- -atrix multiply operations (Figure middle left), which can be efficiently mapped onto hardware (Dao et al., 2022a). The order p controls the number of matrix multiply operations and introduces a tradeoff: higher values of p incur lower FLOP cost via smaller matrices, but require more 1/O communicate intermediate results. Using simple GPU cost model, to adjust p based on the sequence length to balance the FLOP cost and I/O cost. This decomposition introduces a second benefit: reduction the amount of the sequence that needs to be kept in SRAM. which makes kernel fusion viable at longer sequence lengths. As result, FLASHFFTCON V scales across four orders of magnitude in sequence length, from 256 million. FLASHFFTCOM also exploits real-valued FFT algorithm to cut the length of the FFT operation in half (Sorensen et al., 1987), and selectively skips portions the matrix- multtply operations when the input zero-padded.
show how
Finally, the matrix view of the FFT convolution presents natural interface to implement two architectural modifications: partial convolutions, which learn with that is shorter than the input sequence, and frequency-sparse convolutions, which zero out portions of the kernel ky in frequency space. These can be viewed as convolutional analogues to sparsc/approximate attention in Transformers (Han al., 2015a;bb Kitaev et al., 2020; Paliotta et al., 2023; Beltagy et al., 2020), and map naturally on to FLASHFFTCCNN both algorithms can be implemented simply by skipping portions of the matrix decomposition, thus reducing memory footprint and wall-clock runtime.
Evaluation We show that FLASHFFTCONV speeds up the FFT convolution, yieluiin qughtt-quality, more efficient, and longer-sequence models.
2

---

Published as a conference paper at ICLR 2024
Quality FLASHFFTCONY improves the quality of convolutional sequence models via better efficiency: for the same compute budget, FLASHFFTCONY allows Hyena- Paans-GPT-s-- to achieve 2.3 points better perplexity (Poli et al., 2023), and allows BERT-base (Fu et al., 2023a) to achieve up to 3.3 higher average GLUE score- a gain in performance equivalent to doubling the parameters of the model.
Efficiency FLASHFFTCONY makes convolutions more efficient across four orders of magnitude in sequence length, yielding speedups of up to 7.93x and memory savings of up 5.60 over PyTorch. FLASHFFTCONY achieves up to 62.3% end-to-end FLOP utilization- -only 10% less than FlashAttentionn and is faster in wall-clock time than FlashAttention-Vy end-to-end at sequence lengths 2K and longer due to lower FLOP costs.
Longer Sequence Models FLASHFFTCONY enables longer-sequence models. high-resolution image classification, FLASHFFTCON yields the first model that can solve the challenging Path-512 task (sequence length 256K) from the long range arena benchmark (Tay et al., 2020). In DNA modeling, FLASHFFTCON uses partial convolutions to extend HyenaDNA (Nguyen et al., 2023) to 4M sequence length-yielding the first model that can embed the longest human genes (up to 2.3M base pairs) at single nucleotide resolution.
Overall, we hope that FLASHFFTCON enables further adoption of convolutional sequence models and that the insights from our work helps inform the design of better hardware-efficient architectures.
2
BACKGROUND
We provide some background on the FFT convolution and the Monarch FFT decomposition, and discuss the performance characteristics of GPUs.
2.1 FFT CONVOLUTION
Recall the definition of a convolution operation: (u [) = Computing this formula directly incurs O N₂)) FLOPs in sequence length N and kernel lergth Ns. For long convolutions, where N₁ = NN a popular strategy is to use the Fourier transform to convert the signal u and kernel k to the frequency domain, and compute the convolution using pointwise multiplication in frequency domain, using Equation computed in O (NNoogN) time using the FFT- bringiig the overall cost of the long convolution from O(N2) to O(Vlgg)) We note that the FFT convolution technically computes a circular convolution Σn Uy uyki-jj where i j-j <0 loops back to the end of k. For this reason, and k are often padded with zeros to compute causal convolution.
Critically, a Fourier transform Fx over an input of length N can be
Monarch FFT Decomposition For N N2, an order-2 Monarch FFT decomposition rewrites =x=P([N, &FN, DDP-1(N, (IN, &FN, PP, where @ denotes the Kronecker product, .Fx is the N x N permutation matrix that reshapes the input to N, x N2, transposes it to N₂ x Ni, and then reshapes it back to N, and D D x.N is a diagonal matrix containing correctional values called Twiddle factors (Bailey, 1989). Higher-order Monarch decompositions recursively apply the order-2 decomposition to FN, or FN., which reduces FLOP costs but increases the number of permutation operations, increasing I/O cost.
discrete Fourier matrix,
NNN
2.2 GPU PERFORMANCE CHARACTERISTICS
We provide some background on the GPU memory hierarchy and available compute units, well as compute-boung vs. memory-bount operations. We focus on GPU programming in this paper, but the general principles extend to most modern hardware accelerators (Jouppi al., 2023; Zhang et al., 2022a; Lavely, 2022; Emani et al., 2021).
GPU Compute Model and Memory Hierarchy GPUs have a memory hierarchy consisting of high-bandwidth memory (HBM), shared memory (SRAM), and registers, shown in Figure 1 Left. Lower/larger levels the memory hierarchy have more space but are much slower, whereas higher/smaller levels of the memory hierarchy have less space but are much faster (NVIDIA, 2017; 2020; 2022). The memory hierarchy is closely tied to the GPU compute model. GPU composed of many independent streaming multiprocessors (SMs), each of which is composed of independent threads. HBM is shared among all SMs, but each SM has an independent SRAM. The SRAM is shared among all the threads in the SM. Each thread has access to its registers, but cannot access the
3

---

Published as a conference paper at ICLR 2024
registers of other threads. Thus, performing global operations between SMs requires moving data to and from HBM, whereas independent work in each S.M can remain local to SRAM.
GPU Compute Units Modern GPUs (since the V100 (NVIDIA, 2017)) have specialized matrix multiply units called tensor cores, which can compute matrix-mane multiply operations with much higher TFLOPs than the general-purpose compute units. For example, the H100 tensor core can compute matrix multiplication between 16 x 16 matrices at 1.0 PFLOPs, whereas the general-purpose compute units can only compute 67 TFLOPs (NVIDIA, 2022).
Memory-Boune vs. Compute-Boun Operations GPU operations can be memory-bounie or compute- bound. semorry-bound operations are bottlenecked by the amount of I/O between HBM and registers they need to perform, and are limited by the bandwidth of the memory hierarchy. Examples include simple pointwise operations such as addition or multiplication, as well most traditional FFT imple- mentations. Compute- bound operations are bottlenecked by the amount of FLOPs they need to execute, and are limited by the speed of the compute units. Examples include large matrix multiply operations.
Kernel Fusion A popular method for reducing I/O costs is kernel iisioo-loading data for multiple operations into SRAM, computing them independently in each SM, and then writing the final results back to HBM. Kernel fusion is common (and can be automated) for pointwise operations (Paszke al., 2019), but is more challenging for complex operations that require referencing multiple pieces of data. For example, fusing the operations in attention was not common until the development of FlashAttention (Dao al., 2022b).
3 FLASHFFTCONY
Section 3.1 provides a broad overview of FLASHFFCCON and shows how to adapt the Monarch FFT decomposition to convolutions, which involves broadcasting the matrix mul- tiply in parallel across the input sequence. We also describe our kernel fusion strategy and how exploit domain- specific properties of the convolution in ML for further optimization. Section 3.2 presents cost model character- izing the relative cost of different order-p decompositions of the FFT as sequence length changes, along with sim- ple heuristic for selecting p given hardware characteristics. Finally, Section 3.3 discusses architectural extensions by presenting analogues sparsity in convolutional kernels.
F
@
@
Broadcast over B, B.H
Broadcast over N
T.T SRAM
Transpose in SRAM
Replaces Permutations
Figure 2: Top: FLASHFFTCONY adapts the Monarch FFT decomposition to broadcast matrix multiply operations over the sequence instead of over the batch and hidden dimensions. Bottom: This converts HBM permutations into simple matrix transpose operations in SRAM.
3.1 FLASHFFTCONT ALGORITHM
We describe the core FLASHFFTCONY algorithm. Algo- rithm provides an overview. We first describe how we adapt the Monarch FFT decomposition for convolutions. Then, we discuss how the Monarch decomposition enables kernel fusion for long sequences. We conclude by present- ing domain-specific optimizations.
Adapting Monarch for Fusion The Monarch FFT decomposition, as well as classical algorithms such as Bailey's 's FFT algorithm (Bailey, 1989), traditionally broadcasts the matrix operation against the batch dimension and the hidden dimension (Figure 2 top left). This allows each FNz operation in the I,, FFN, matrix to run independently. However, it also makes kernel fusion difficult; fusing across the matrix multiply and permutation operations requires loading at least 16 sequences at once into SRAM to fill out the matrix multiply unit unit-limiting sequence length to around 2 on A 100 and H100.
Instead, we broadcast the matrix operation across the entire sequence (Figure 2 top right) and run the algorithm in parallel across the batch and hidden dimensions. This reduces the SRAM requirements for kernel fusion, since we only need to load a single sequence into SRAM at a time aelowwng us to fuse the entire kernel for sequences up to 32K in fp16/bf16 on A100 and H100. Broadcasting along the sequence has an added benefit: the permutations simply become matrix transposes (Figure bottom), which can be done quickly using well-established routines on-chip (NVIDIA, 2020). The core algorithm is shown in Algorithm for two-way decomposition. Higher-order decompositions and more details are given in Appendix B.
4

---

Published as a conference paper at ICLR 2024
Algorithm FLASUFFFCCNN core algorithm, with order-2 Monarch decomposition. We assume N N=N for simplicity here.
Input: Input R convolution kernel y C FFT matrices F E CNNN,, F p-1
C Twiddle factors E CN.. tinv E CNN, tile size Brates H tile size Hate-
Output: Output BxHxx
for SMs in parallel across B/ Buite x H Hate do
Load F. F fime from HBM.
for 11-111110 Hrue do
Load K fotkkyii
ky [h] from HBM, reshaped to N1 x N1.
for b+ to Biste do
Load X 4- ubbhh from HBM, reshaped N1 x N
4 ( X) F
D FFT, decomposed into two steps
X - K+
Elementwise multiply with ky
< ((XF-1)T + kmm) F-1
b Inverse FFT, decomposed into two steps
Write Y to HBM.
Cost Model for Order-p Monarch Decompositions
Cost Per Token C/N
100
Cores, p
10
Matricas Too
SRAM Limit
p
p=
p
256
1K
4K
16K
64K
256K
1M
4M
Sequence Length
Figure Compute costs of different order-p Monarch decompositions as sequence length increases on A100. For short sequences, higher-orde decompositions break the FFT down into matrices that are too small for tensor cores. length 32K, the sequence too long to fit into SRAM, and higher-order decompositions are necessary for fusion.
Kernel Fusion and Recomputation The Monarch decomposition allows kernel fusion for long sequences. Inner layers of the decomposition do not require the entire sequence, which reduces the SRAM requirements for fusion. Thus, long sequences, we can fuse the innermost matrix operations and elementwise multiplications, and take an I/O each for the outermost matrix operations. We use also use recomputation in the backward pass to reduce the memory footprint and I/O cost. Instead of storing intermediate results HBM for the backward pass (e.g., the intermediate result of Fxu), we simply recompute them the backward pass.
Domain-Specine Optimizations Finally, we use few domain-specifer optimizations to adapt the convolution specifically for the sequence learning workload. First, since the convolutions used in sequence learning real-to-real convolutions (with real kernel weights), called one-stage decimation in time to compute the FFT of a sequence length N using complex FFT of length N/2 (see Appendix B)-cctting the FFT cost in half. Second, inputs and outputs are often padded with zeros in the convolution to compute a causal convolution (Gu et al., 2021; ; Poli et al., 2023; Fu et al., 2023a). We special-case this padding, and it to eliminate half of the outermost matrix multiply operations in the FFT and iFFT. We also fuse in additional operations around the convolution, such as elementwise-gatie gating. to further reduce I/O.
a classic algorithm
3.2
CoST MODEL ORDER-P MONARCH DECOMPOSITION
We present a formal cost model for an order-p Monarch decomposition of the convolution based on sequence length. The cost model accounts for both the cost of compute and I/O, similar to roofline analysis (Hennessy & Patterson, 2011) 1). Let and H be the batch size and model hidden dimension, respectively, and assume that compute the convolution in half precision. Let N be the sequence
5

---

Published as a conference paper at ICLR 2024
length, and let N = - NY be the product of factors. For simplicity, will assume that N is power of Let μ be the size of the matrix-matrin multiply unit on the GP (e.g. 16 for A 110 (NVIDIA, 2020) and H100 (NVIDIA, 2022)). Let To and Tar be the empirically-achieeee for general-purpont define ((N)) helper function that returns TG if N, < μ., and TM if N; μ.. Finally, let a and as be aeppircaall-achievable bandwidth for HBM and SRAM, respectively. Sample values for these constants are given Appendix D.
FLOPs on the GPU arithmetic, and matrix-matrix multiply arithmetic,, respectively. For convenience,
Now, we can present the cost of an FFT convolution with an order-p Monarch decomposition. Let (i) be helper function that returns z the intermediate result of step i is stored on HBM, or os otherwise. The cost of the convolution using an order-p Monarch decomposition is given by:
16NN,
4N
C=BH
(2)
(N.)
w(i)
Figure 3 graphs Equation 2 for different order-p decompositions on different sequence lengths for A100, for pe (244 For cases where N1 =... = Np. the total FLOP cost an order-p decomposition grows with ((N())))))) However, for shorter sequences, higher-order decompositions are actually more expensive, since they decompose to matrices that are smaller than the maarix-matril multiply unit (corresponding to the early bumps). Note also the bump in cost for p p= between 32K and 64K, which is result of running out of SRAM but which is mediated by extra decomposition for p=4.
3.3
ARCHITECTURAL EXTENSIONS: SPARSITY IN CONVOLUTIONS
We present two architectural extensions to FLASHFFTCONY: partial convolutions and frequency- sparse convolutions, which are approximate convolution methods enabled by FLASHFFTCON CoNv.
partial convolutions, we zero out later portions of the convolution kernel, analogous to local attention. This has two benefits. First, it reduces the memory footprint, since it requires fewer elements to be held in GPU memory at once. Second, allows for natural extensions of pretrained convolutional model to longer sequences (i.e., via sliding window approach).
Partial Convolutions
Prequency-Sparye Convolutions In frequency-spant convolutions, we zero out portions of the convolution kernel in frequency space, i.c. zeroing out portions of ky. Here, the specific sparsity pattern can yield computational benefits-coe correspodirggg to skipping portions of the matrix multiplies in FLASHFFTCONV (see Appendix B for examples).
4 EXPERIMENTS
In this section. we evaluate FLASHFFTCON in terms of quality and efficiency. First (Section 4. 4.1), we show that FLASHFFTCONY allows models to achieve better quality for the same compute budget in language modeling-mmtchinng the performance of models with twice the parameters for free. FLASHFFTCONV also enables higher quality via higher resolution in image classification-see -solving the challenging Path-512 512 task for the first time simply via increased sequence length. Next (Section 4.2), we demonstrate FLASuFFTCON
speedup for convolutions, evaluate its efficiency gains end-to-end, to Transformers using FlashAttention-v2.
and compare convolutional model using Finally (Section 4.3), we evaluate partial and frequency-spe -paasee convolutions. Partial convolutions yield the first DNA model that can embed the longest genes at single nucleotide resolution (2.3M base pairs), and frequency-sparse convolutions yield speedup while maintaining- or improving-qualll
4.1 IMPACT OF EFFICIENCY ON QUALITY
We study how FLASHFFTCONY impacts downstream quality. First, given two implementations with the same compute budget, FLASHFFTCONY achieves higher quality due higher training throughput. Second, we show that improved efficiency can lead to higher quality via longer sequence length.
Improvement in Quality with Fixed Compute Budget To evaluate the impacts of efficiency on downstream quality, we train two popular convolutional language models, M2-BERT-bbas. (Fu et al., 2023a) and Hyena-s (Poli et al., 2023), from scratch. These models are trained BERT-style (masked language modeling) and GPT-style (next token prediction), respectively. We compare the quality of models trained with the same compute budget but different implementations of the convolution either FLASHFFT CONV Py PoTch implementation of the FFT convolution. FLASHFFTCONY achieves higher pretraining throughput, which allows the models to see more data during pretraining. These efficiency gains improve average GLUE score by up 3.4 points for M2-BERT-bass and perplexity

---

Published as a conference paper at ICLR 2024
Table 1: Improvement in quality given fixed compute budget.
Model (Metric)
PyTorch
FLASHFFTCONV
M2- 22-BERT-basse 110M (GLUE Score 1)
77.6
80.9
1555na-s-155M (PPL 4)
13.4
11.1
Table 2: Classification accuracy (f) on Path-X and Path-512 from the long range arena benchmark (Tay et al. 2020). FLASHFFTCONV allows for higher-resolution classification. X indicates out of memory.
Task (seq. len.)
PyTorch FLASHFFTCONV
Path-X (16K)
96.9
96.9
Path-512 (256K)
X
96.1
Table 3: Time (L) to compute the forward pass of convolution with FLASHFFTCONY in milliseconds on one H100-SXM, as well as ablations removing specific optimizations. We also show memory savings. Results scaled to batch size 64, hidden dimension 768.
P
p=3
p
Sequence Length
256
1K
4K
8K
16K
32K
1M
2M
4M
PyTorch
0.43
1.57
6.65
13.7
28.6
62.1
2,346.3
4,892.1
10,127.6 7,587.0
CLASSFFTCOON
0.09
0.24
1.37
3.19
21.8
1,492.8
0.21
0.67
3.51
7.71
21.4
45.5
Speedup over Py Porch
4.78 x
6.54x
4.85x
x.29
3.09 x
2,85
1.57
1.82x
x
Memory Savings
8.21x
7.73x
7
7,21x
2.64X
2.63x
2.63
by 2.3 points for Hyena-s. These improvements in quality are similar in magnitude to the effect of doubling the number of parameters in the model (see Appendix C for reference results).
Longer Sequence Models Next, we show how increased efficiency can lead to higher quality via longer sequence lengths. We evaluate long convolution models on Path-X and Path-512, high-resolution imaging tasks from the long range arena (LRA) benchmark (Tay et al., 2020)." These tasks take an image (128x x 128 for Path-X and 512x512 for Path-512), flatten out, and require sequence model to classify whether two dots in the image are connected by path.
Existing PyTorch implementations of convolutional sequence models (or even prior optimized imple- mentations (Fu et al., 2023b)) fail to achieve better-than-rame thad-m (50%) accuracy on Path-512 due to out of memory errors and a lack of support for such long sequences. Table 2 shows that FLASH FLASHFFTCONY allows a convolutional sequence model to solve Path-512 2 for the first time simply by increasing the available sequence length and reducing the memory footprint of the model through fusion.
4.2 EFFICIENCY
We evaluate FLASHFFTCONY on how fast computes convolutions compared to PyTorch and how much speedup it yields for convolutional sequence models end-to-end. We also evaluate memory savings compared to PyTorch and compare against Transformers using FlashAttention-22 (Dao, 2023).
We benchmark the speed of the convo- lution compared against an FFT convolution implemented in Torch. We also benchmark ablations evaluating kernel fusion without using tensor cores which recovers the strong bascline of using Nvidia's 's cuFFTdx kernel fusion library (NVIDIA, 2023a) and FLASHFFTÇONY without its domain- specific optimizations.
FLASHFFT FCCONV Provides Speedup and Memory Savings
Table shows that FLASHFFTCONY outperforms PyTorch FFT convolution across all sequence lengths, by up to 6.54x. Appendix C shows further speedup. up to 7.93 for domain-specific optimizations. Speedups are greatest for short sequences, where the PyTorc FFT convolution is dominated by 1/O costs. Speedup is more modest for longer sequences, which incur additional I/O costs (between registers and SRAM for the p= 3 and between SRAM and HBM for 4) Without using
We refer Path-512 scaled-up version of Path-256.

---

Published as a conference paper at ICLR 2024
Table 4: End-to-end throughput (1) of convolutional sequence models against PyTorch.
Model (size, seqlen, unit)
PyTorch
FlashFFTConv
Speedup
12-BERT-baase (110M, 128, seqs/s)
4,480
8,580
1.9x
Hyena-s-4K (155M, 4K, seqs/s)
84.1
147
1.7x
Long convs, Path-X (102M, 16K, images/s)
126
308
2,4x
SaShiMi (5.4M, 64K, audio clips/s)
38.7
50.3
1.3x
HyenaDNA (IM, seqs/s)
0.69
3.03
4.4x
Table 5; End-to-end throughput (1) in thousands of tokens per second, FLOP utilization, and speedup of Hyena against GPT running FlashAttention-v2 (Dao, 2023) across sequence lengths for 100..
Model
2K
8K
16K
GPT-2.7B, FA-v2 (Dao, 2023)
33.8
27.8
21.6
Hyena-2.7B, FLASHFFTCONY
35.2
35.2
32.3
FA-v2 FLOP Utilization
65.7
72.1
78.5
FLASHFFTCONY FLOP Utilization
62.3
61.9
56.5
FLASHFFTCCNN Speedup
1.1x
1.3x
1.5x
the Monarch decomposition for tensor cores (fusion-only), FLASHFFTCONV becomes bottlenecked by the speed of general arithmetic operations on GPUs and the size of SRAM. Table also shows memory savings over PyTorch. FLASHFFTCONV reduces the memory footprint of convolutions via recomputation in the backward pass and kernel fusion.
FLASHFFTCONT Speeds Up Convolutional Sequence Models We benchmark end-to-end through- put of convolutional sequence models across various modalities and sequence lengths spanning four orders of magnitude. We benchmark M2-BERT (Fu et al., 2023a), BERT-style language model that has sequence length 128; Hyena-s-4K (Poli et al., 2023), a GPT--ttyle language model with se- quence length 4K; covoouuooolllll model (Fu al., 2023c) trained on Path-X with sequence length 16K (Tay et al., 2020); SaShiMi (Goel et al., 2022), an audio generation model trained on 1-second audio clips sampled at 64 KHz; and HyenaDNA- IM (Nguyen et al., 2023), a DNA modeling model trained on IM sequence length. Details of the architectures and ceattecture-speciffe optimizations (such as fusing multiplicative gating for M2 and Hyena models) are given in Appendix D.
Table 4 shows that FLASHFFTCONV speeds up these models end-to-end. Speedup varies vary by the size of the models and the relative amount of time spent computing the convolution compared to other parts of the models. FLASHFFTCONV only speeds up the SaShiMi model by 1.3x, since the model interleaves convolutions with SSM- filter generation, pooling layers, and MIPs, which reduces the relative amount of time spent computing the convolution itself. Speedup is greatest for HyenaDNA, where Py PyTocch is bottlenecked by small batch size. The PyTorch implementation only allows batch size 1 on 80GB GPU, whereas FLASHFFTCONY allows batch size 4 yielding significant speedup.
MASHFFFCONN is Faster than 2lastAttention-22 We compare end-to-end efficiency of a 2.7B Brr--araeeter Hyena model using FLASHFFTCONY against 2.7B ..7B-paramettr GPT model us- ing Ftaetontet (Dao, 2023) at three sequence lengths. Table 5 shows throughput, end-to-end FLOP utilization, and speedup. FLASHFFTÇONY achieves lower end-to-end FLOP utilization than FlashAttention-222 but achieves higher throughput, since convolutions incur fewer overall FLOPs.
4.3 PARTIAL AND FREQUENCY-SPAPEE CONVOLUTIONS
We evaluate the impact of partial convolutions on downstream quality and memory footprint and on how well they can extend the sequence length of existing models. We evaluate the impact of frequency- sparse convolutions downstream quality, and we show that rrequency-sparte convolutions can yield up to 1.4x additional speedup in the convolution without impacting quality.
Partial Convolutions Reduce Memory Footprint and Increase Sequence Length Partial convolu- tions reduce the memory footprint of models, in both language modeling and DNA modeling. A large proportion of the convolution filters can be pruned without impacting downstream quality. Table
8

---

Published as a conference paper at ICLR 2024
Table 6: Quality and memory footprint of partial convolutions during training across sequence lengths.
Hyena-s-8K
8K
4K
2K
1K
512
256
PPL (U)
13.8
13.8
13.8
13.9
14.0
14.2
Memory Footprint (U)
32.5G
15.3G
11.8G
8.4G
6.1G
5.8G
Table 7: PPL (4) from using partial convolutions extend the sequence length of HyenaDNA longer sequences. At 4M sequence length, the models are able to embed the longest human genes.
Base Filter Length
1M
2M
4M
HyenaDNA-4500
2.91
2.91
2.91
HyenaDNA-IM
2.91
2.91
2.90
Table 8: Applying frequencyy-paarsiy to the filters of a pretrained HlyenaDNA-Ia model.
Sparsity Fraction
0%
50
75%
79 %
84%
91 %
PPL. (4)
2.91
2.91
2.90
2.91
2.93
2.98
Convolution Speedup (f)
1.0x
1.2x
1.3x
1.4x
1.5x
1.8x
shows that a Hyena-s-8K model can be pretrained with much shorter convolution kernel-as short as 2 negatively impacting quality.
Partial convolutions yield another benefit: we can naturally extend the sequence length of existing pretrained models. We extend a pretrained HyenaDNA-1M, model to 4M sequence length with promising results (Table 7) --yieeldngg the first model that can embed the longest human genes at single-nuclcotide resolution (2.3M base pairs) (See Appendix C for visualization of gene embeddings).
Trequency-Spars Convolutions Increase Throughput Frequency-sparse convolutions can increase the speed of convolutions and may also have positive effects on quality. Table 8 shows that we can set up to 79% of the entries of the kernel F to zero without losing quality. Sparsification in frequency space may even improve the quality of pretrained models slightly; ; the PPL of pretrained HyenaDNA-IM model improves by 0.01 points after its kernels are 75% sparsified in frequency space- potentiltl result of removing tigh-frequency noise. Sparsification also yields up to 1.4x speedup in the convolution via skipping entire blocks of the matrix-matrin multiplies in the Monarch decomposition. Appendix D provides more details about the sparsity patterns used in Table 8.
5 RELATED WORK
We provide an abbreviated summary of related work. An extended related work can be found in Appendix A. Long convolutional models have emerged as promising alternative to Transformers for sequence modeling (Gu et al., 2021; Romero et al., 2021b; lb; Poli et al., 2023; Ma et al., 2022; Hasani et al., 2022; Smith al., 2023), partieularly for for long-context applications such as DNA modeling and speech synthesis [gguyen et al., 2022; Goel et al., 2022). Our work is related long history of efficient FFT algorithms, ranging from Cooley-Tukey (Cooley & Tukey, 1965) to more recent FFT algorithms (Ayinala al., 201 1; Chu & George, 1999; Bahn et al., 2009; Bailey, 1989). Our work builds and extends these algorithms to adapt them to convolutions. Finally, partial and frequency- convolutions relate long history of sparsity in deep learning. from pruning (Han et al., 201 2015a;b; Sanh et al., 2020) to sparse structured matrices (De Sa et al., 2018 8; Sindhwani et al., 2015).
CONCLUSION
We present FLASHFFTCON new system for optimizing FFT convolutions for long sequences. We show that FLASHFFTCONY improves quality under fixed compute budget, enables longer-sequence models, and improves the efficiency of long convolutions. We also show that analogues of sparsity in convolution filters map naturally to FLASHFFTCONY compute model, and reduce memory footprint and runtime. We hope that our work will help support further adoption of convolutional sequence models, and that insights can help inform the design of future architectures.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
We gratefully acknowledge the support of NIH under No. US4EB020405 (Mobilize), NSF under Nos. CCF2247015 (HardwareAAware) CCF1763315 (Beyond Sparsity), CCF1563078 (Volume to Velocity), and 1937301 (RTML); US DEVCOM ARL under Nos. W911NF-23-2-01 84 (Long-context) and W911NF-21-2-0251 (Interactive Human- AI Teaming); ONR under Nos. N0001423 1263 (Deep Stanford HAI under No. 247183; NXP, Xilinx, LETI-CEA, Intel, IBM, Microsoft, NEC, Toshiba, TSMC, ARM, Hitachi, BASF, Accenture, Ericsson, Qualcomm, Analog Devices, Google Cloud, Salesforce, Total, the HAI-GCP Cloud Credits for Research program, the Stanford Data Science Initiative (SDST), and members of the Stanford DAWN project: Meta, Google, and VMWare. The U.S. Government is authorized reproduce and distribute reprints for Governmental purposes notwithstanding any copyright notation thereon. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views, policies, or endorsements, either expressed or implied, NIH, ONR, or the U.S. Government.
Signal Processing);
REFERENCES
Gustaf Ahdritz, Nazim Bouatta, Sachin Kadyan, Qinghui Xia, William Gerecke, Timothy O' D'nnelll, Daniel Berenberg, lan Fisk, Niccolò Zanichelli, Bo Zhang, et al. Openfold: Retraining alphafold2 yields new insights into its learning mechanisms and capacity for generalization. bioRxiv, pp. 2022
Ben Athiwaratkun, Sujan Kumar Gonugondla. Sanjay Krishna Gouda, Haifeng Qian, Hantian Ding, Qing Sun, Jun Wang, Liangfu Chen, Jiacheng Guo, Parminder Bhatia, et al. On io-efficient attention mechanisms: Context-aware bifurcated attention and the generalized multi group attention. In Workshop on Efficient Systems for Foundation Models@ ICML2023, 2023.
Ziga Avsec, Vikram Agarwal, Daniel Visentin, Joseph R Ledsam, Agnieszka Trabska-Barwinnse Kyle R Taylor, Yannis Assael, John Jumper, Pushmeet Kohli, and David R Kelley. Effective gene expression prediction from sequence by integrating long-range interactions. Nature methods, 18 (10):1196-12033 2021.
Manohar Ayinala, Michael Brown, and Keshab K Parhi. Pipelined parallel fft architectures via folding transformation. IEEE Transactions on Very Large Scale Integration (VLSI) Systems, 20(6): 1068-1081, 2011.
Jun Ho Bahn, Jung Sook Yang, Wen-Hsiang Hu, and Nader Bagherzadeh. Parallel fft algorithms network-on-chips. Journal Circuits, Systems, and Computers, 18(02):255-269, 2009.
David H Bailey. Ffts in external of hierarchical memory. In Proceedings of the 1989 ACM/IEEE conference on Supercomputing, pp. 234 1989.
AJAA Bekele. Cooley-tukey fft algorithms. Advanced algorithms, 2016.
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05/50. 0.00000.5555.... 2020.
Alberto Bietti and Julien Mairal. Invariance and stability of deep convolutional representations.
Advances in neural information processing systems, 30, 2017.
Mikolaj Binkowski, Jeff Donahue, Sander Dieleman, Aidan Clark, Erich Elsen, Norman Casagrande, Luis C Cobo, and Karen Simonyan. High fidelity speech synthesis with adversarial networks. In International Conference on Learning Representattonns, 2019.
Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportunities and risks of foundation models. arXiv preprint arXiv:2108.07258, 2021.
E Oran Brigham. The fast Fourier transform and its applications. Prentice-Hall, Inc., 1988.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing processin systems, 33:1877--1011 2020.
10

---

Published as a conference paper at ICLR 2024
Beidi Chen, Tri Dao, Kaizhao Liang, Jiaming Yang, Zhao Song, Atri Rudra, and Christopher Ré. Pixelated butterfly: Simple and efficient sparse training for neural network models. 2021a.
Beidi Chen, Tri Dao, Eric Winsor, Zhao Song, Atri Rudra, and Christopher Ré. Scatterbrain: Unifying sparse and low-rank attention. In Advances in Neural Information Processing Systems (NeurIPS), 2021b.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009. 14794, 2020.
Aakanksha Chowdhery, Sharan Narang. Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.0233 2022.
Eleanor Chu and Alan George. Inside the FFT black box: serial and parallel fast Fourier transform algorithms. CRC press, 1999.
James W Cooley and John W Tukey. An algorithm for the machine calculation of complex fourier series. Mathematics of computation, 19(90):29-330, 1965.
Fiona Cunningham, James Allen, Jamie Allen, Jorge Alvarez-Jarreta, Ridwan Amode, Irina M Armean, Olanrewaju Austine-Orinan Orimoloyo, Andrey G Azov, If Barnes, Ruth Bennett, et al. Ensembl 2022. Nucleic acids reseurch, (ODDI):D988-DD) D995, 2022.
Zihang Dai, Guokun Lai, Yiming Yang, and Quoe Le. transforranfoormer Filtering out sequential redundancy for efficient language processing. Advances in neural information processing systems, 53:4471-42833 2020.
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning, arXiv preprint vrXiv:2307.08691, 2023.
Tri Dao, Albert Gu, Matthew Eichhorn, Atri Rudra, and Christopher Ré. Learning fast algorithms for linear transforms using butterfly factorizations. arXiv preprint / 00..0.0......... 2020.
Tri Dao, Nimit S. Sohoni, Albert Gu, Matthew Eichhorn, Amit Blonder, Megan I Leszczynski, Atri Rudra, and Christopher Ré. Kaleidoscope: An efficient, learnable representation for all structured linear maps. arXiv preprint arXiv:2012. 2021.
Tri Dao, Beidi Chen, Nimit S Sohoni, Arjun Desai, Michael Poli, Jessica Grogan, Alexander Liu, Aniruddh Rao, Atri Rudra, and Christopher Ré. Monarch: Expressive structured matrices for efficient and accurate training. International Conference on Machine Learning. PMLR, 2022a.
Tri Dao, Daniel Y. Fu. Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-efficient exact attention with aw-awaressss. Advances in Neural Information Processing Systems, 2022b.
Shmuel Bar David, Itamar Zimerman, Eliya Nachmani, and Lior Wolf. Decision s4: Efficient sequence-based rl via state spaces layers. In The Eleventh International Conference on Learning Representations, 2022.
Christopher De Sa, Albert Gu, Rohan Puttagunta, Christopher Ré, and Atri Rudra. A two-pronged progress in structured dense matrix vector multiplication. In Proceedings of the Twenty-Ninth Annual ACM-SIAM Symposium on Discrete Algorithmss, pp. 1060-1079. SIAM, 2018.
Tim Dettmers and Luke Zettlemoyer. Sparse networks from scratch: Faster training without losing performance. arXiv preprint arXiv: 1907,04440 2019.
Prafulla Dhariwal, Heewoo Jun, Christine Payne, Jong Wook Kim, Alec Radford, and Ilya Sutskever. Jukebox: A generative model for music. arXiv preprint urXiv:2005.00:: 2020.
Sander Dieleman, Aaron van den Oord, and Karen Simonyan. The challenge of realistic music generation: modelling raw audio at scale, Advances in neural information processing systems, 31, 2018.
11

---

Published as a conference paper at ICLR 2024
Xin Dong. Shangyu Chen, and Sinno Pan. Learning to prune deep neural networks via layer-wise optimal brain surgeon. Advances in Neural Information Processing Systems, 30, 2017.
Nan Du, Yanping Huang, Andrew M Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, et al. Glam: Efficient scaling of language models with mixture-of-experts. In International Conference Machine Learning, pp. 5547-5569. PMLR, 2022.
Yuli Eidelman and Israel Gohberg. On new class of structured matrices. Integral Equations and Operator Theory, 34(3):293-324, 1999.
Murali Emani, Venkatram Vishwanath, Corey Adams, Michael E Papka, Rick Stevens, Laura Florescu, Sumti Jairath, William Liu, Tejas Nama, and Arvind Sujeeth. Accelerating scientific applications with sambanova reconfigurable dataflow architecture. Computing in Science & Engineering, 23(2): 114-119, 2021.
Yassir Fathullah, Chunyang Wu, Yuan Shangguan, Junteng Jia, Wenhan Xiong, Jay Mahadeokar, Chunxi Liu, Yangyang Shi, Ozlem Kalinli, Mike Seltzer, et al. Multi-head state space model for speech recognition. arXiv preprint arXiv:2305. 12498, 2023.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. The Journal of Machine Learning Research, 23(1): $232-5270, 2022.
Quentin Fournier, Gaétan Marceau Caron, and Daniel Aloise. practical survey on faster and lighter transformers. ACM Computing Surveys, 2021.
Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. arXiv preprint arXiv: 1803.03635, 2018.
Jonathan Frankle, Gintare Karolina Dziugaite, Daniel Roy, and Michael Carbin. Stabilizing the lottery ticket hypothesis. arXiv preprint arXiv: 1tt.:/99330111... 2019.
Jonathan Frankle, Gintare Karolina Dziugaite, Daniel Roy, and Michael Carbin. Linear mode con- nectivity and the lottery ticket hypothesis. In International Conference on Machine Learning, pp. PMLR, 2020.
Daniel Y. Fu, Simran Arora, Jessica Grogan, Isys Johnson, Sabri Eyuboglu, Armin W. Thomas, Benjamin Spector, Michael Poli, Atri Rudra, and Christopher Ré. Monarch Mixer: A simple sub-quadratic GEMM-based architecture. In Advances in Neural Information Processing Systems, 2023a.
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. International Conference on Learning Representations, 2023b.
Daniel Fu, Elliot L. Epstein, Eric Nguyen, Armin W. Thomas, Michael Zhang, Tri Dao, Atri Rudra, and Christopher Ré. Simple hardware-efffcient long convolutions for sequence modeling. International Conference on Machine Learning, 2023c.
Karan Goel, Albert Gu, Chris Donahue, and Christopher Ré. It's raw! audio generation with state-space models. arXiv preprint arXiv:2202.097229 2022.
Albert Gu, Karan Goel, and Christopher Re. Efficiently modeling long sequences with structured state spaces. In International Conference on Learning Representations, 2021.
Albert Gu, Ankit Gupta, Karan Goel, and Christopher Ré. On the parameterization and initialization of diagonal state space models. In Advances in Neural Information Processing Systems, 2022a.
Albert Gu, Isys Johnson, Aman Timalsina, Atri Rudra, and Christopher Ré. How train your hippo: State space models with generalized orthogonal basis projections. arXiv preprint arxiv:2006. 12037, 2022b.
12

---

Published as a conference paper at ICLR 2024
Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces effective structured state spaces. In Advances Neural Information Processing Systems, 2022.
Song Han, Huizi Mao, and William J Dally. Deep compression: Compressing deep neural networks with pruning, trained quantization and huffmmn coding. arXiv preprint wr..v:1510000110100001 20100011911111111 2015a.
Song Han, Jeff Pool, John Tran, and William Dally. Learning both weights and connections for efficient neural network. Advances in neural information processing systems, 28, 2015b.
Ramin Hasani, Mathias I Lechner, Tsun-Huang Wang, Makram Chahine, Alexander Amini, and Daniela Rus. Liquid structural state-space models. arXiv preprint arXiv:2209. 12951, 2022.
Kaiming He, Xiangyu Zhang. Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
John L Hennessy and David A Patterson. Computer architecture: a qquntitative approach. Elsevier, 2011.
Md Mohaiminul Islam, Mahmudul Hasan, Kishan Shamsundar Athrey, Tony Braskich, and Gedas Bertasius. Efficient movie scene detection using state-space transformers. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 18749-18758, 2023.
Yanrong Ji, Zhihan Zhou, Han Liu, and Ramana V Davuluri. DNABERT:
pre-trained bidirectional encoder representations from transformers model for DNA-language in genome. Bioinformatics, 37 (15):2112-2120, 2021.
Norm Jouppi, George Kurian, Sheng Li, Peter Ma, Rahul Nagarajan, Lifeng Nai, Nishant Patil, Suvinay Subramanian, Andy Swing. Brian Towles, et al. Tpu v4: An optically reconfigurable supercomputer for machine learning with hardware support for embeddings. In Proceedings of the 50th Annual International Symposium on Computer Architecture, pp. 1-14, 2023.
Thomas Kailath, Sun-Yuan Kung, and Martin Morf. Displacement ranks of matrices and linear equations. Journal of Mathematical Analysis and Applications, 68(2):395 407, 1979.
Nal Kalchbrenner, Erich Elsen, Karen Simonyan, Seb Noury, Norman Casagrande, Edward Lockhart, Florian Stimberg, Aaron Oord, Sander Dieleman, and Koray Kavukcuoglu. Efficient neural audio synthesis. In International Conference Machine Learning, pp. 2410 2410-2419. PMLR, 2018.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mnns: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning, pp. 5156-5165. PMLR, 2020.
Sanghyeon Kim and Eunbyung Park. Smpconv: Self-moving point representations for continuous con- volution. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 10289-10299, 2023.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint arXiv:2001. 044551, 2020.
Alex Krizhevsky, Ily Sutskever, and Geoffrey Hinton. Imagenet classification with deep convolu- tional neural networks. Advances in neural information processing systems, 25, 2012.
Kundan Kumar, Rithesh Kumar, Thibault De Boissiere, Lucas Gestin, Wei Zhen Teoh, Jose Sotelo, Alexandre De Brebisson, Yoshua Bengio, and Aaron Courville. Melgan: Generative adversarial networks for conditional waveform synthesis. Advances in neural information processing systems, 32, 2019.
Mitsuru Kusumoto, Takuya Inoue, Gentaro Watanabe, Takuya Akiba, and Masanori Koyama. A graph theoretic framework of recomputation algorithmss for memory-efficiem backpropagation. Advances in Neural Information Processing Systems, 32, 2019.
13

---

Published as a conference paper at ICLR 2024
Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng. Lianmin Zheng, Cody Hao Yu, Joseph E. Gonzalez, Hao Zhang. and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the ACM SIGOPS 29th Symposium Operating Systems Principles, 2023.
Kushal Lakhotia, Eugene Kharitonov, Wei-Ning Hsu, Yossi Adi, Adam Polyak, Benjamin Bolte, Tu-Anh Nguyen, Jade Copet, Alexei Baevski, Abdelrahman Mohamed, el. On generative spoken language modeling from raw audio. Transactions of the Association for Computational Linguistics, 9:1336-1354, 2021.
Adam Lavely. Powering extreme-scale hpc with cerebras wafer-scale accelerators. Cerebras White Paper, 2022.
Binrui Li, Shenggan Cheng, and James Lin. tefft: Accelerating half- salf-precision fft through tensor cores. arXiv preprint arXi::2104. 14471, 2021.
Raymond Li, Loubna Ben Allal, Yangtian Zi, Niklas Muennighoff, Denis Kocetkov, Chenghao Mou, Mare Marone, Christopher Akiki, Jia Li, Jenny Chim, et al. Starcoder: may the source be with you! arXiv preprint arXiv:2305.06 2023.
Yuhong Li, Tianle Cai, Yi Zhang, Deming Chen, and Debadeepta Dey. What makes convolutional models great on long sequence modeling? arXiv preprint arXiv:2210.09298. 2022.
Lin, Yongming Rao, Jiwen Lu, and Jie Zhou. Runtime neural pruning. Advances in neural information processing systems, 30, 2017.
Hao Liu and Pieter Abbeel. Blockwise parallel transformer for long context large models. arXiv preprint arXiv: 19370, 2023.
Zhuang Liu, Hanzi Mao, Chao- Yuan Wu, Christoph Feichtenhofer, Trevor Darrell, and Saining Xie. A convnet for the 2020s. In Proceedings of the IEEE/CVE conference on computer vision and pattern recognition, pp. 11976-1988, 2022.
Xuezhe Ma. Xiang Kong, Sinong Wang, Chunting Zhou, Jonathan May, Hao Ma, and Luke Zettlemoyer. Luna: Linear unified nested attention. Advances in Neural Information P Processing Systems, 34: 1441-2455 2021.
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.10653 2022.
Temesgen Mchari and Nils Strodthoff. Towards quantitative precision for ecg analysis: Leveraging state space models, self-supervision and patient metadata. IEEE Journal of Biomedical and Health Informatics, 2023.
Harsh Mchta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. Long range language modeling via gated state spaces, arXiv preprint ar0iv: /3947, 2022.
Maxim Milakov and Natalia Gimelshein. Online normalizer calculation for softmax. arXiv preprint arXiv: /805.02867 2018.
Koichi Miyazaki, Masato Murata, and Tomoki Koriyama. Structured state space decoder for speech recognition and synthesis. In ICASSP 2023-2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 1-5. IEEE, 2023.
Eric Nguyen, Karan Goel, Albert Gu, Gordon Downs, Preey Shah, Tri Dao. Stephen Baccus, and Christopher Ré. S4nd: Modeling images and videos as multidimensional signals with state spaces. In Advances in neural information processing systems. 2022.
Eric Nguyen, Michael Poli, Marjan Faizi, Armin Thomas, Callum Birch- Birch-Sykes, Michael Wornow, Aman Patel, Clayton Rabideau, Stefano Massaroli, Yoshua Bengio, et al. Hyenadna: Long-rane genomic sequence modeling at single nucleotide resolution. In Advances in Neural Information Processing Systems, 2023.
14

---

Published as a conference paper at ICLR 2024
NVIDIA. Nvidia Tesla V100 GPU architecture, 2017.
NVIDIA. Nvidia A100 tensor core GPU architecture, 2020.
NVIDIA. Nvidia H100 tensor core GPU architecture, 2022.
NVIDIA. cuffidx v1. documentation, 2023a. http:://lossssaddaacccccoouulllt
NVIDIA. Cutlass 3.2, 2023b. https//gtttubbcoom/////AAAAAsss
NVIDIA. Cuda c++ programming guide, 2023c. http:///osssniddacoooccuudddeeep guide/index.html
Aaron van den Oord, Sander Dieleman, Heiga Zen, Karen Simonyan, Oriol Vinyals, Alex Graves, Nal Kalchbrenner, Andrew Senior, and Koray Kavukcuoglu. Wavenet: A generative model for raw audio. arXiv preprint arXiv: 1609.0396 2016.
Alan Oppenheim. Applications of digital signal processing. Englewood Cliffs, 1978.
Alan V O'ppenheim, John R Buck, and Ronald W Schafer. Discrete-tine signal processing. Vol. 2. Upper Saddle River, NJ: Prentice Hall, 2001.
Daniele Paliotta, Matteo Pagliardini, Martin Jaggi, and François Fleuret. Fast causal attention with dynamic sparsity. In Workshop on Efficient Systems for Foundation Models@ ICML2023, 2023.
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. Proceedings of the 40th International Conference on Machine Learning (ICML 2023), 2023.
Ryan Prenger, Rafael Valle, and Bryan Catanzaro. Waveglow: A flow-based generative network for speech synthesis. In ICASSP 2019-2019 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 3617-3621. IEBE, 2019.
David W Romero, Robert-Jan Bruintjes, Jakub M Tomczak, Erik J Bekkers, Mark Hoogendoorn, and Jan van Gemert. Flexconv: Continuous kernel convolutions with differentiable kernel sizes. arXiv preprint arXiv:21 10.08059,
David W Romero, Anna Kuzina, Erik Bekkers, Jakub Mikolaj Tomczak, and Mark Hoogendoomn. Ckconv: Continuous kernel convolution for sequential data. In International Conference on Learning Representations, 2021b. 20211.
Victor Sanh, Thomas Wolf, and Alexander Rush. Movement pruning: Adaptive sparsity by fine-tuning. Advances in Neural Information Processing Systems, 33:20378-20388, 2020.
Jonathan Shen, Ruoming Pang, Ron J Weiss, Mike Schuster, Navdeep Jaitly, Zongheng Yang. Zhifeng Chen, Yu Zhang, Yuxuan Wang, Rj Skerrv- et al. Natural uts synthesis by conditioning wavenet on mel spectrogram predictions. 2018 IEEE international conference on acoustics, speech and signal processing (ICASSP),
IEEE, 2018.
Vikas Sindhwani, Tara Sainath, and Sanjiv Kumar. Structured transforms for soall-footprint deep learning. Advances in Neural Information Processing Systems, 28, 2015.
Jimmy TH Smith, Andrew Warrington, and Scott Linderman. Simplified state space layers for sequence modeling. In The Eleventh International Conference on Learning Representations, 2023.
V Sorensen, Jones, Michael Heideman, and Burrus. Real-valued fast fourier transform algorithms. IEEE Transactions acoustics, speech, and signal processing, 35(6):849-86) 1987.
15

---

Published as a conference paper at ICLR 2024
Siyi Tang, Jared A Dunnmon, Liangqiong Qu, Khaled K Saab, Christopher Lee-Messer, and Daniel L Rubin. Spatiotemporal modeling of multivariate signals with graph neural networks and structured state space models. arXiv preprint arXiv:2211.11177 2022.
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena: A benchmark for efficient transformers. In International Conference on Learning Representations, 2020.
Yi Tay, Mostafa Dehghani, Jai Prakash Gupta, Vamsi Aribandi, Dara Bahri, Zhen Qin, and Donald Metzler. Are pretrained convolutions better than pretrained transformers? In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing ( (Volume I: Long Papers), pp. 4349 4349-4359, 202 la.
Yi Tay, Vinh Tran, Sebastian Ruder, Jai Gupta, Hyung Won Chung, Dara Bahri, Zhen Qin, Simon Baumgartner, Cong Yu, and Donald Metzler. Charformer: Fast character transformers via gradient- based subword tokenization, arXiv preprint arXiv:2106. 12672, 2021b.
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: survey. ACM Computing Surveys, 55(6))1-28, 2022.
Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. Advances in neural information processing systems, 30, 2017.
Jue Wang, Wentao Zhu, Pichao Wang, Xiang Yu, Linda Liu, Mohamed Omar, and Raffay Hamid. Selective atattuuedd for long-form video understanding. Propeee Proceedings of the IÉFF/CVF Conference on Computer Vision and Pattern Recognition, pp. 6387-6397, 2023.
Junxiong Wang, Jing Nathan Yan, Albert Gu, and Alexander M Rush. Pretraining without attention. arXiv preprint arXiv:22 12. 10544, 2022a.
Qipeng Wang, Mengwei Xu, Chao Jin, Xinran Dong, Jinliang Yuan, Xin Jin, Gang Huang, Yunxin Liu, and Xuanzhe Liu. Melon: Breaking the memory wall for resource-efficient on devicee machine learn- the ing. In Proceedings the 20th Anmual International Conference on Mobile Systems, Applications
and Services, pp. 450-463, 2022b.
Sinong Wang, Belinda Zii, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006. 04768, 2020.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Yang, et al. Big bird: Transformers for longer sequences. Advances in neural information processing systems, 13:17223-17299 2020.
Dan Zhang, Safeen Huda, Ebrahim Songhori, Kartik Prabhu, Quoc Le, Anna Goldie, and Azalia Mirhoseini. A full-stack search technique for domain optimized deep learning accelerators. In Proceedings of the 27th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, pp. 27-42, 2022a.
Michael Zhang, Khaled Kamal Saab, Michael Poli, Tri Dao, Karan Goel, and Christopher Re. Ef- fectively modeling time series with simple discrete state spaces. In International Conference on Learning Representations, 2022b.
Chen Zhu, Wei Ping, Chaowei Xiao, Mohammad Shoeybi, Tom Goldstein, Anima Anandkumar, and Bryan Catanzaro. Long-short transformer: Efficient transformers for language and vision. Advances in Neural Information Processing Systems,
2021.
16

---

Published as a conference paper at ICLR 2024
APPENDIX
We presenteetendeeeee extended related work (Appendix A), additional algorithmic details (Appendix B), additional experimental results (Appendix C), and experimental details (Appendix D).
A EXTENDED RELATED WORK
Long Convolutions in Sequence Modeling Long convolutional models have emerged as a promising alternative to Transformers for sequence modeling (Gu et al., 2002a;b;; Romero et al., 2021 b;a; Poli et all, 2023; Ma et al., 2022; Fu et al., b0; Nguyen et al., 2023; Hasani et al., 2022; Smith et al.,, 2023). These methods differ in how they generate the convolutional kernels; for example. the S4 line of work uses learned state space models (Gu et al., 2021; Mehta et al., 2022; Gupta et al., 2022; Ma al., 2022), while other works (Poli et al., 2023; Romero et al., 202 parameterize the convolution using an MLP from positional encodings. However, all the models operate by taking convolution over the input sequence with kernel as long as the input: y = u* where E TR. BxxxXxX N & R HxNN and the kernel & is broadcast along the B dimension. When used for language modeling, these models often incorporate elementwise multiplicative gating as well: y= f()) @ ((g()))h(e)) @h(u))*)) where f. g. and h are linear maps along the !l dimension (Poli et al., 2023; Fu et al., 2023b;a; Mehta et al., 2022; Wang et al., 2022a).
Long-Context Applications Long convolutional models have especially been helpful for long- context applications, such as DNA modeling and speech synthesis. In DNA modeling, most longer- context genomic models have relied on either tokenization (Ji et al., 2021; Zaheer et al., 2020; Tay et al., 20) or downsampling (Fournier et al., 2021; Avsec et al., 2021). However, recent work has suggested that modeling DNA directly from base pairs can yield downstream improvements in quality, which requires long sequence lengths (Nguyen et al., 2023).
Like DNA modeling, speech synthesis has also benefited from long-contest modeling. While traditional speech synthesis pipelines use intermediate representations such as spectrograms (Kumar et al., 2019; Prenger et al., 2019; Shen et al., 2018), 1), linguistic features (Binkkwski et al., 2019; Kalchbrenner et al., 2018; Oord et al., 2016). or discrete audio codes (Dhariwal et al., 2020; Dieleman et al., 2018; Lakhotia et al., 2021; Van Den Oord et al., 2017), recent work has shown that modeling the speech directly from the raw waveform can yield downstream improvements in quality (Goel et al., 2022). Again, such models require long sequences to model audio at the rate at which it is naturally sampled, necessitating long-sequence modeling.
FFT Algorithms There is a long history of efficient FFT algorithms, ranging from the Cooley-Tukey FFT algorithm published in 1965 (Cooley & Tukey, 1965) to parallel FFT algorithms (Ayinala et al., 2011) and more (Chu & George, 1999; Bahn et al., 2009; Bailey, 1989). These algorithms have enabled fundamental progress in range of disciplines, from control theory (Brigham, 1988; Bekele, 2016) to signal processing (Oppenheim, 1978; Oppenheim et al., 2001). As FFTs prove more useful for modern deep learning applications, such as long convolutions, new techniques are required to run them efficiently modern accelerators. Our work continues line of work exploring how to use tensor cores for the FFT convolution (Fu et al., b023c;b; Li et al., 2021), and extends the algorithmic capabilities to much longer sequences.
Sparsity in Deep Learning As deep learning models have grown larger and deeper (Bommasani al., 2021; Brown et al., 2020; Chowdhery et al., 2022), there is increasing interest in reducing the cost of training and running models. Sparsity in particular has received a great deal of attention, and has a long history in machine learning, including work in pruning neural networks (Han et al., 2015a;b; Sanh et al., 2020; Lin et al., 2017; Dong et al., 2017) and finding lottery tickets (Frankle & Carbin, 2018; Frankle et al., 2019; 2020). Our work in partial convolutions and frequency-sparse convolutions relates to this line of work, as an analogue of sparsity in convolutional filters. The Monarch decomposition is also closely related to structured matrices. Structured matrices have subquadratic (o(n²) for dimension X n)) parameters and runtime, such as sparse and low-rank matrices, and fast transforms (Fourier, Chebyshev, sine/cosine, orthogonal polynomials) (Dao et al., 2022a). Structured matrices can often be computed with simple divide-and-onn and--conueerrrrreeeeeeeeeeeae--e schemes, and can be used to represent many fast transforms (De Sa et al., 2018; Sindhwani et al., 2015; Kailath et al., 1979; Eidelman & Gohberg, 1999).
17

---

Published as a conference paper at ICLR 2024
Optimization of deep learning primitives There is rich history of optimizing deep learning primitives. Many techniques, such as kernel fusion, aim to reduce data movement. Recently, libraries such PyTorch 2.0 Paszke et al. (2019) have added kernel fusion automatically. Other techniques include checkpointing. wherein one stores fewer intermediate results and recomputes the others on-the- fly where they are needed, trading additional compute for memory Kusumoto al. (2019); Wang et al. (2022b). Many algorithms also have ontiiiiotisztiions that can remove unnecessary computation or memory accesses Milakov & Gimelshein (2018).
Another line of optimization techniques aims to reduce FLOPs. MLPs and attention are particularly popular targets of FLOP reduction, via sparse factorizations of weights (Cooley & Tukey, 1965; Frankle & Carbin, 2018; Dao et al., 2022a; 2020; Zhu et al., 2021; Dao et al., 2021; Dettmers & Zettlemoyer, 2019; Chen et al., 2021a), or sparse/low-rank approximations of attention Beltagy al. (2020); Kitaev et al. (2020); Ma et al. 1) Fedus et al. (2022); Du et al. (2022); Wang et al. (2020); Dai et al. (2020); Zhu (2021); Choromanski al., 2021b; Tay et al., 2022).
(2020); Katharopoulos et al. (2020) and their combinations (Chen
B ALGORITHM DETAILS
B.1 DOMAIN-SPECIRE OPTIMIZATIONS
We review the details of how to compute a real-to-real FFT of size N using complex FFT of size N/2, following tutorial by (Sorensen et al., 1987).
For this section, we adopt notation common in describing FFT algorithms. Let x(n) be an input sequence of length N, and let X (k) be the result of its discrete Fourier transform. Recall that:
X(k) Σ r(n)WT*
(3)
for l: = 0,1,..,,NNN 1, where WN = 2xi/N is the Nth root of unity.
First, if (n)) is real, then symmetries emerge in X (k). In particular, we have x X (N k),, where denotes complex conjugation. These symmetries allow us to have an algorithm for computing X (A)) using single complex DFT of size N/2.
In particular:
((k) Σ
W
n=0
N/2-1
N/2-1 W Σ n=0
= Σ
xt (2(n)
x(2n+1)WWM)
R=0
for k= 0. 1. The DFT is now decomposed into two parts: a DFT over the even-indexed elements of (n), and over the odd-indexed elements of (n)).
We can now create a third complex sequence, length N/2, and put the even-indexed elements of a=n) in the real part, and the odd-indexed elements of (n) in the imaginary part. Let:
=(n) for n n= N,,,,,.,N//211 Then, we compute the N/2-sized DFT (kk), and we can recover the DFT over the even and odd parts of (n) (X. [k] and X, respectivell
X. ( 2-
Iokk] -
Zik] - Z IN/2-
2i
We can now recover = using:
[ mod N/2]+ *x.kk mod N/2
The inverse FFT proceeds similarly. The goal is to recover (n)) given input X(k), using simple complex inverse DFT of length N/2.
First, recover X. [k] and X, [k):
18

---

Published as a conference paper at ICLR 2024
X.[k: X xX]]/X"[NN/
X
(N/N/2-
X,[A=]=
W
for 0/2 Then, we construct [[::
=|k]=X, -6]
We the inverse DFT to recover 2(n), and then recover z=(n) from the real and imaginary parts of =(n):
x(2n) Re(zn) x (2n+1) Im (zn);
for n- 0/2---,.,,//11
To implement these in our kernels, we perform the bookkeeping after reading the inputs or before writing the output, and then use the FFT/FFFT implementations detailed Algorithm and others.
B.2 LOW-LEVEL CUDA DETAILS
To ensure high performance, we implement CUDA kernels for each specific sequence length, allowing us to cater to specific performance nuances that arise from the decomposition at that sequence length. In this section, we dive into some of the low-level implementation details for FLASHFFTCONV.
Matrix Multiplication Using CUDA Tensor cores CUDA Tensor cores can perform the multi- plication of two k and k matrices for bfloat16 floatl6 elements, using around the same number of cycles as is required for the multiplication two scalars. m X k n must be of one the following: 16 x 16 x 16, 32 x 8 x 16, 8x32x 16. This informs choice of radix for decomposition when performing the FFT and iFFT. In particular our implementation breaks down matrix-matrix multiplications into blocked matrix-matrix multiplications where k x 16 16 16. We note the following about matrix-matrix multiplication on tensor cores (NVIDIA, 2023c):
Tensor cores are utilized at the level of the warp and programmatic access of the tensor cores is via the Warp Level Matrix Multiply Accumulate (WMMA) API.
Tensor core operands held in register fragments (wmmma :: matrix a, and wmmna :: matria_l) Ji and results are written to a register fragment (wmmma accumalalorr..
The operand fragments can hold data in row-major or column-major format and data in the wmma:: :: accumulator fragment can be written memory in row-major or column-major format.
The specific mapping of items in fragment to threads in warp is unspecified, however, the mapping of items to threads in the wmma:: accumuulator fragment exactly matches that for the wmmaa: :: matriz. aa fragment read row-major, allowing us to directly copy the results of matrix-matro multiplication and use as the operand for another matrix-matrix multiply.
To perform a matrix matrix multiplication C A B using the tensor cores, a warp loads the contents of A and B into registers (WMMA fragments in CUDA parlance), performs the matrix matrix multiplication, and writes the results which are stored in an accumulator fragment back to memory.
key part of ensuring high performance is minimizing I/O across different levels of the memory hierarchy: from HBM to SRAM and from SRAM to registers. To ensure this, we move the output from the accurrulator fragment directly into matriz_o _a fragment for use subsequent matrix multiplications, avoiding extra trip to SRAM. However, this is only possible if the output from the previous matrix-matrii multiply does not need to be transposed before using it as an operand for the next one. When this is not the case, we need to make trip to SRAM and back. In Algorithm 2 detail I/O from SRAM to registers.
Register Reuse
Locality and Tiling The algorithm is trivially parallelizable across B and H, allowing us to tile in both dimensions at the threadblock level. In Algorithm all loops from <- to N₁ are warp-tiled.
19

---

Published as a conference paper at ICLR 2024
Algorithm 2 Detailed Annotation of FLASHFFTCONN core algorithm showing I/O from SRAM register fragments, with two-way Monarch decomposition. We assume N for simplicity here.
Input: Input #
convolution kernel y C FFT matrices F E CNIX,, F C Twiddle factors E CN
tinv CaavC" ECN, tile size Beates H tile size Hade-
Output: Output E y B
for SMs in parallel across B/ Buite x H Hase do
Load F, F-
fime from HBM.
for h+- to Hrse do
Load K k [h] from HBM, reshaped to N1 X1.
for b+ 1 to Biste do
Load X [ from HBM, reshaped N1 x NN.
4 F
> F (matrix a), X (malrix b) output to accumulator
Load X from accumulator
matria_a
EElementwise multiply directly in matria X (matrix.a) F (matrix b) output accumulator Load X from accumulator to matrix_a
X + X
XF
x X*Kr
b Elementwise multiply with ky directly in matria a X (malrix_a), F-1 (malrix__) output accumulator
X XF-1
Write X from accumulator fragment to SRAM
Load X from SRAM to matria a fragment
X X stino
> Elementwise multiply with tine directly in motria X (matrir.a). F-1 (matrix:)) b) output accumulator
Y-XF-1
Write Y to HBM.
Miscellaneous optimizations In addition to the above optimizations, we also perform some other optimizations that provide marginal speedup. These include: utilizing vector intrinsics/types for performing memory reads/writes and arithmetic for 16-bit floating point (fp16) and brain float point (bf16), allowing non-tensor speed. Furthermore, we double buffer I/O movements across all levels of the memory hierarchy, reducing warp stalls. We also aggressively tune our kernel hyperparameters such block and tile dimensions, and loop unrolling factors for the best performance on the specific underlying hardware.
operations these types be performed at around twice the normal
B.3 GENERALIZATION TO 3-WAY AND 4-Y MONARCH DECOMPOSITIONS
We provide algorithm listings for 3-way and 4-way Monarch Decompositions.
3- Bay Decomposition Algorithm 3 shows the algorithm for a 3-wayy Monarch decomposition. involves one extra matrix multiply operation either side of the FFT and iFFT, and proceeds over the algorithm in Algorithm in an inner loop.
4-way Decomposition For the 4-way decomposition, we assume that we need write intermediate outputs to HBM. Here, treat the 3-way decomposition sub-routine, and assume it has a fused kernel (i.e., Algorithm 3). We compute one matrix multiply for the FFT and one for the iFFT, and then call the kernel for the 3-way decomposition over the rows of the output. The algorithm is listed in Algorithm 4.
B.4 FREQUENCY-SPAREE PATTERNS
We describe frequency-sparse patterns and the matmul savings in more detail here. We use the full w-w decomposition case, since the algorithms generalize to lower-order decompositions.
Let N= Nf. and consider kemee ky E C Consider the matrix multiply and looping operations that occur when computing the FFT portions of FLASHFFTCONY (u, ky) (the iFFT portions are the same, in the opposite order):
1. In Algorithm 4, there is one FFT operation over the columns of u, reshaped to N₁ N/N1, and a Twiddle correction..
2. Then, Algorithm 3 iterates the rows of for a:= N₁ steps.
20

---

Published as a conference paper at ICLR 2024
Algorithm 3 algorithm for 3-way decopppstion... assume N N2 for simplicity here.
Input:
Input RB convolution kernel K CCHN FFT matrices F E C F
CNNN,, Twiddle factors 1 E C.V,, f1,in0 E Of* N1
12 ECE 12 EC*, B tile size Brites H tile
size Hrite
Output: Output IR
for SMs in parallel across B/ Biile X H/HHstee do
Load F, F.F t, time from HBM.
for h Hrate
Load K < kg from HBM, reshaped to N2 N1.
K₁ - K7-
> Transpose last two dimensions.
Reshape K to N1 N.. for b to Biste do
Load x [bbh) from HBM, reshaped N1 NN X
for sori+11 N2 do
X' 1 N1: (++1))
X[:i+N N1+1)
b Transpose, matmul, transpose.
X Xxte for 1 1 N₁ do
b> LLop over rows
X'+FFXI
Reshape X' to N N1
X'+ ((F X' F
FFT, decomposed into two steps Elementwise multiply with ky
X'+X'* X'
Y' Y' ((X'F 1)T "*lino)
> Inverse FFT, decomposed into two steps
Y'-Y"
Y
>> Finish inner loop
Y Y-Yxt2,,m
for it to N do Y ( N [i+1)+N] Y|i (i+1)
b Transpose, matmul, transpose.
Write Y to HBM.
3. Let u' be the row in a specific iteration. In Algorithm 3, there an FFT over the columns of u', reshaped N, x N77 and Twiddle correction.
4. Then, the inner loop iterates over rows u' for 3: Nj steps.
5. In each loop, u' has one FFT operation with a twiddle factor correction. Let the matrix of this FFT operation be denoted A.
6. Then there a second FFT operation. Let the matrix of this FFT operation denoted B.
Now, reshape ky to N1 x Nı x N1 x N1. Let us consider how sparsity along the each of the four dimensions of k lets skip operations in the above steps.
Sparsity in the first dimension allows us to skip computation in B, exactly in proportion to how much of the first dimension eliminate. This can result in cost savings, long as B can still be expressed using the tensor cores on-chip after skipping the computation. For example, if B is 32x32, then N1 32, and does not make sense eliminate more than half of the first dimension.
Sparsity in the second dimension works exactly the same way, except it allows to skip computation in A.
Sparsity in the third dimension lets us reduce 3. Each row of the third dimension that we remove lets us skip one iteration of the inner loop in step 4 above.
Sparsity the fourth dimension lets us reduce a. Each row of the fourth dimension that we remove lets skip one iteration of the outer loop in step 2 above.
21

---

Published as a conference paper at ICLR 2024
Algorithm 4 4 o here.
algorithm for 4ay decompostiinnn We N Ni for simplicity
Input:
Input E | convolution kernel k CC*N FFT matrices F E C F-1 E
CNIXN, Twiddle factors E C. tinv E E CN CN t2.mm CCY.
Output: Output TR Bx
Reshape u to B H x X (N)N)).
Reshape y H x XNx (N/N)).
ky
ky kj.
Transpose last two dimensions.
Reshape ky H N/N1
u u Fu
Computes the FFT over the columns of u.
Reshape u to B (H (N/NI).
> Move Nj into dimension.
Reshape y to (HN1) (N/M)
Call FLASHFFTCON (u, k)).
Call 3-way FLASHFFTCONV.
Reshape to B X y+FF 'u Return y.
P Computes the iFFT over the columns of u.
Table 9: Sparsity patterns for k and sparsity fraction for the frequency-sparte convolution experiment Table
Sparsity Pattern
S
==0,b=0,c=0,d=0
a= =66.b=0,c0,0000
50
== 16,b=16,c=0
75
a= 1=16bb=16.c44.d=1
79
a= 0=1,,b=66,08,,55
84
a=16.b=11,.
91
Table 10: Time (1) to compute the forward pass of gated convolution with FLASHFFTCONV in milliseconds on one H100-SXM. We also show memory savings. All results scaled batch size 64, hidden dimension 768. p indicates the order of the Monarch decomposition.
p
P
3
p
Soquence Length
256
4K
8K
16K
32K
IM
4M
PyTorch
29.9
3,071.4
FLASHFFTCONV
0.11
0.29
1.43
3.58
12.2
26.3
1,768.9
4,623.5
10,049.4
Speedup
5.64x
7.93x
6.64x
5.42x
2,45x
3,22
1.74%
1.37
1.30
Momory Savings
6.65x
6.40x
6.35x
6.34x
6.17x
5.87x
2,82x
2.81x
2
As example, reveal the sparsity dimensions that applied in the experiment detailed in Table 8 in the main paper. Conceptually, we use the full 2-million length kernel k and reshape it to 32 x 32 x 32 x64 64. Let Specifically, we set ky :....... = 0. ky [..]]...]] =0, 0. ky e.:..:.:] :::]== 0, and k [:....ddd sequentially. The formula the sparsity fraction S given a. in this case is given
c, and d be variables describing how much of each dimension we zero.
5 (32 a) (32 (32 (2)))))(644222 (64-
or more generally, 1 minus the product of the fraction of each dimension that is removed. Table lists the configurations of the sparsity patterns and the sparsity fractions used for the experiment in Table 8.
B.5 HARDWARE SUPPORT
FLASHFFTCON developed on A100 GPUs, and tested on A100 and H100 GPUs. Older generations of GPU such as V100 are not supported, since the sizes of the tensor cores are different. We look forward to integrating more general libraries such as Cutlass (NVIDIA, 2023b) support wider range of GPUs, and developing support for non-GPU accelerators.
C ADDITIONAL RESULTS
22

---

Published as a conference paper at ICLR 2024
Table 11: Full results for the forward pass of convolution with FLASHFFTTONV compared to
PyTorch in milliseconds on one H100-SXM. Batch size 64, hidden dimension 768.
Seq Len
PyTorch
FLASHFFTCONV
Speedup
256
0.43
0.09
4.69
512
0.81
0.15
5.34
1024
1.57
0.24
6.61
2048
3.27
0.55
5.95
4096
6,65
1.37
4.87
8192
13.72
3.19
4.30
16384
28.58
9.27
3.09
32768
62.09
21.84
2.84
65536
141.15
67.96
2.08
131072
292.26
147.26
1.98
262144
582.76
308.48
1.89
524288
1,167.28
742.26
1.57
1048576
2,346.26
1,492.84
1.57
2097152
4,892.09
2,695.51
1.81
4194304
10,127.56
7,586.96
1.33
Table 12: Full results for the forward pass of gated convolution with FLASHFFTCONV compared to PyTorch in milliscconds on one H100-SXM. Batch size 64, hidden dimension 768.
Seq Len
PyTTrrh
FLASHFFTCONY
Speedup
256
0.62
0.11
5.76
512
1.18
0.19
6.14
1024
2.30
0.29
7.81
2048
4.70
0.67
7.05
4096
9.49
1.43
6.65
8192
19.38
3.58
5.42
16384
39.91
12.18
3,28
32768
84.79
26.32
3.22
65536
186.69
79.84
2.34
131072
382.98
181.51
2.11
262144
764.08
376.96
2.03
524288
1,530.34
878.93
1.74
1048576
3,071.37
1,768.94
1.74
2097152
6,342.58
4,623.46
1.37
4194304
13,031.21
10,049.42
1.30
C.1 SPEEDUP FROM DOMAIN- -PECCICCC OPTIMIZATIONS
We benchmark domain specific optimizations in FLASHFFTCONT Table 10 shows the performance of a gated convolution v ((w@w) k). where and are linear projections of the input This pattern is common convolutional and 5SM-based architectures for language modeling (Fu et al., 2023c;b; Poli et al., 2023; Mehta al., 2022). A Py'Tochh implementation of gated convolution incurs additional I/O overhead from the gating operations, whereas FLASUFFTCONY fuses the gating operations into the convolution. This fusion results in further speedup over Py PyTorch, up to 7.93x.
C.2 FULL RESULTS FOR ALL SEQUENCE LENGTHS AND SETTINGS
We report full results for all sequence lengths in powers of two between 256 and 4M. We report full results for five cases:
Table 11: Standard forward pass, where the FFT size is the same the input size. This is equivalent to a circular convolution.
Table 12: Gated forward pass, where the FFT size is the the input size.
23

---

Published as a conference paper at ICLR 2024
Table 13: Full results for the forward pass of convolution where the input is half the length of the convolution size with FLASHFFTCONV compared to PyTorch in milliseconds on one H100-SXM. Batch size 64, hidden dimension 768.
Seq Len
PyTorch
FLASHFFTCONY
Speedup
256
0.44
0.09
4.64
512
0.82
0.16
5.03
1024
1.57
0.24
6.45
2048
3,25
0.53
6.08
4096
6.59
1.37
4.83
8192
13.60
3.13
4.34
16384
28.37
8.82
3.22
32768
61.87
21.34
2.90
65536
141.42
77.32
1.83
131072
292.26
151.28
1.93
262144
582.82
315.99
1.84
524288
1,167.21
757.33
1.54
1048576
2,343.55
1,525.13
1.54
2097152
4,922.63
3,321.71
1.48
4194304
10,179.86
7,305.61
1.39
Table 14: Full results for the forward pass of a gated convolution where the input is half the length of the convolution size with FLASHFFTÇON compared to PyTorch in milliseconds on one H100-SXM. Batch size 64, hidden dimension 768.
Seq Len
PyTorch
FLASHFFTCONY
Speedup
256
0.54
0.11
4.71
512
1.01
0.19
5.27
1024
1.94
0,29
6.75
2048
3.97
0.59
6.69
4096
8.01
1.41
5.68
8192
16.42
3.46
4.75
16384
34.04
10.62
3.21
32768
73.15
25.03
2.92
65536
163.75
78,88
2.08
131072
337.37
153.13
2.20
262144
672.48
319.47
2.10
524288
1,346.99
763.97
1.76
1048576
2,704.91
1,538.89
1.76
2097152
5,644.20
3,545.79
1.59
4194304
11,625.79
8.132.32
1.43
Table 13: Forward pass, where the input size is half the FFT size. This is equivalent to causal convolution.
Table 14: Gated forward pass, where the input size is half the FFT size.
Table Standard backward pass, where the FFT size is the same as the input size.
Table Memory use for FLASHFFTCONY compared to PyTorch for a convolution, scaled to batch size 64, hidden dimension 768.
Table 17 Memory for gated convolution using FLASHFFTCONN compared to PyTorch for a convolution, scaled to batch size 64, hidden dimension 768.
Speedups vary, but generally follow the trend from the results in the body of the paper. FLASHFFT- CONY achieves significant memory savings over Porch due to recomputation in the backward pass and kernel fusion. To measure memory savings, we measure the relative additional memory from calling the convolution operations (we do not measure the footprint of hte original inputs).
24

---

Published as a conference paper at ICLR 2024
Table 15: Full results for the backward pass of convolution with FLASHFFFTCON compared to
PyTorch in milliseconds on one H100-SXM. Batch size 64, hidden dimension 768.
Seq Len
PyTorch
FLASHFFTCONY
Speedup
256
0.76
0.24
3.24
512
1,45
0.22
6.43
1024
2.83
0.65
4.37
2048
5.76
1.48
3.90
4096
11.56
2.86
4.05
8192
23.11
6.16
3.75
16384
46.85
18.57
2.52
32768
103.85
57.68
1.80
65536
241.81
111.76
2.16
131072
489.38
239.32
2.04
262144
976.24
519.49
1.88
524288
1,960.31
1,240.95
1.58
1048576
3,938.92
2,708.36
1.45
2097152
7,909.27
4,977.93
1.59
4194304
16,552.21
12,932.02
1.28
Table 16: Memory usage in GB for FLASHFFTCONV compared to PyTorch. Scaled up to batch size
64, hidden dimension 768.
Seq Len
PyTorch
FLASHFFTCONY Memory Reduction
256
0.42
0.05
8.21 x
512
0.80
0.10
8.19x
1024
1.58
0.20
7.73x
2048
3.12
0.39
7.94
4096
6.21
0,82
7.61 x
8192
12.39
1.63
7.59
16384
24.93
3,46
7.21x
32768
50.43
7.68
6,57x
65536
121.60
46.08
2.64x
131072
243.21
92.18
2,64x
262144
486.41
184.39
2.64x
524288
972.83
368.91
2,64x
1048576
1945.65
738.34
2.64x
2097152
3889.23
1477.69
2.63x
4194304
7778.45
2961.56
2.63x
C.3 REFERENCE LARGER MODELS
Table 18 gives performance numbers for larger models trained for the same number of tokens and steps as the reference Py PyTorch models in Table in the main paper.
The GPT-style PyTorch models are trained for 5B tokens, with batch size 512K tokens. The BERT- style PyTorch models are trained for 16000 steps, with batch size 64K tokens. In contrast, the FLASHFFTCONY models, with higher training throughput, are trained for 15B tokens and 70000 steps in the same compute budget, respectively.
C.4 DNA EMBEDDINGS
We use our 4M-seuuence length HyenaDNA model to generate embeddings for various DNA seg- ments following the procedure from (Nguyen al., 2023). The DNA classes include human genes corresponding to different biological function annotations from the Ensembl genome dataset known as biotypes (Cunningham al., 2022). The longest human gene, the dystrophin gene, is annotated.
25

---

Published as a conference paper at ICLR 2024
Table 17: Memory usage in GB for FLASHFFTCONV for gated convolution compared to Py PyTorch.
Scaled up to batch size 64, hidden dimension 768.
Seq Len
PyTorch
FLASHFFTCONY
Memory Reduction
256
0.66
0.10
6.65x
512
1.28
0.19
6.61 x
1024
2.54
0.40
6.40
2048
5.04
0.78
6,49
4096
10.05
1,58
6.35x
8192
20.07
3.17
6.34
16384
40.29
6.53
6,17x
32768
81.15
13.83
5.87x
65536
164.61
58.37
2.82x
131072
329.22
116.75
2.82x
262144
658.44
233.54
2,82x
524288
1316.89
467.21
2.82x
1048576
2633.78
934.95
2.82x
2097152
5265.48
1870.90
2881x
4194304
10530.97
3747.99
2.81x
Table 18: Reference quality numbers for models when trained for the same number of steps and training data.
Model (Metric)
M2-BERT -aase--base-ee 110M (GLUE Score f)
77.6
M2-BERT-large-26M (GLUE Score f)
81.0
Hyena-s- 155M (PPL !)
13.4
Hyena-m 3y355- (PPL !)
11.1
Table 19: Too error for FLASHFFTCCON compared to PyTorch FFT convolution in fp32 across sequence lengths, for fp16 and bf16.
256
1K
4K
32K
1M
fp16
3,8c-6
7.6e-6
1.5e-5
6.5e-5
4.0e-4
bf16
3.1e-5
3.1e-5
6.1e-5
6.1e-5
6.
Table 20: Time spent in the convolution operation for the models in Table 4, on H100.
Model % Time in Convolution
Speedup
M2-BERT
48.5%
1.9x
38.4% Long Convs
Hyena
1.7x
64.3%
2.4x
SaShiMi
22.1%
1.3x
HyenaDNA
81.8%
4.4x
C.5 PRECISION OF FLASHFFTCONY
FLASHIFTCOON operates over inputs in fp16 6 bfl6, which enables faster execution (tensor cores 2x faster than fp32 matrix multiply) at the cost of lower-precision computation. We measure the error from computing an FFT convolution using FLASHFFTCONV compared to computing it in PyTorch in fp32 precision. Table 19 shows the results, using an Lx norm. Error increases with sequence length, but is still small even for sequences up 1M.
26

---

Published as a conference paper at ICLR 2024
Unprocessed Pseudogene
miRNA
IncRNA
Protein Coding
Processed Pseudogene
SnRNA
Miscellaneous RNA
TEC
SnoRNA
Dystrophin Gene
Longest Human Gene Dystrophin, 2.3M Base Pairs
Figure 4: t-SNE visualization of various genes and DNA segments using our new HyenaDNA-4M, The longest human gene, Dystrophin, is annotated.
C.6 PERCENTAGE OF TIME SPENT IN CONVOLUTION
Table 4 shows different amounts of speedup for different models end-to-end. Speedup varies by the amount of time each model spends computing the convolution, compared to other operations like MLP and convolution kernel generation.
Table 20 presents the amount of time PyTorch implementation of each model spends computing convolutions. We profile each model on H100 using the PyTorch CUDA profiler and manually inspect the amount of time spent the convolution. Speedup closely correlated with how much time the PyTorch implementation spends in convolution.
D EXPERIMENT DETAILS
D.1 COMPUTE
All experiments were conducted on box with 8xA100-40GB GPUs box with 8xH100-SXM GPUs.
D.2
FIXED COMPUTE BUDGET EXPERIMENT
For the experiment in Table 1, we train an M2--ERT-barl model from scratch, and Hyena-s- 155M model from scratch.
We train the 12-BERT-base model using masked language modeling of 30% on the C4 dataset, and fine-tune it on GLUE using the protocol from (Fu al., 2023a). The FLASHFFTCONY model has higher training throughput, so it trains for more tokens; we train the FLASHFFTCONY model for 70,000 steps with batch size 64K tokens. The PyTorch model, with lower training throughput, only trains for 16,000 steps, with the same batch size. The M2-BERT-bass. model we use is marcheter-matched with a Transformer BERT-base. expansion factor of four. It also uses block-diagonal MLP with four blocks. The M2 Hyena filter has embedding dimension 5, filter order 128, and initial sine activation factor of 10. We train with learning rate 8e-4, weight decay le-5, and 6% warmup with cosine decay.
has 12 hidden layers, with model dimension of 960, and an
We train the Hyena-s- 155M model using a causal language modeling objective on the Pile. We train the FLASHFFT CONV model for 15M tokens, and the PyTorch model for 5M tokens. The Hyena-s- 155M model matches the configuration from (Poli et al., 2023) and has 18 layers, with hidden dimension of 864, and expansion factor of 4, The Hyena filter has embedding dimension 33, filter order 64, and
27

---

Published as a conference paper at ICLR 2024
initial sine activation factor of 14. We train with learning rate 6e-4, with 1% warmup time and a cosine decay.
D.3 PATH-X AND PATH-512 EXPERIMENTS
For the experiment in Table 2, we use simple convolutional language models, as in (Fu et al., 2023c).
For Path-X, we use the same model and hyperparameters' the convolutional model from (Fu et al., 2023c). We use a convolutional model with 6 layers. prenorm batch norm, and hidden dimension of 256. For the convolution filter parameters, we use kernel dropout 0.3, kernel learning rate 0.0005, factor 0,001, and two channels on the filter. We use an overall learning rate of 0.0005 and weight decay 0.05. We train for 500000 steps, with I 10000 steps of warmup with a cosine decay, and global batch size 16.
For Path-512, we scale up the resolution of Path-256. We train for 200000 steps, with 10000 steps warmup, learning rate 0.0005, and weight decay 0.05. For the model, we train with 4 layers, and hidden dimension 256. We use kernel dropout 0.1, kernel learning rate 0.0005, factor 0.001, and two channels on the filter. We keep the filter length to be 65536.
D.4 CONVOLUTION BENCHMARKS
For the experiments in Table 3, we time the forward pass of convolution with batch size 64, hidden dimension 768, and varying sequence length. If we run out of memory for a sequence length, we split the batch and hidden dimension and call the forward pass multiple times. We time each call 30 times and take the average of the runs. We use the same protocol for the backward pass in Table ??.
D.5 END-TO-ENT END MODELING DETAILS
For the experiments in Table 4, run forward pass of each model, and use it to compute throughput. Batch sizes vary by model, and we check throughput calculations with a few batch sizes to make sure the result is consistent. For the M2-BERT-bass model, et al., 2023a). For the Hyena-s-444 4-KK model, we use an identical model to the one in Table 1, but with filter length of 4K. For the long convs Path-X model, we use the same model as in Table 2. For the SaShiMi model, we use the standalone SaShiMi model from the official implementation (Goel et al., 2022), and we use 8 layers with hidden dimension 64, and 4 up pool and down pool layers. For the HyenaDNA model, we use the official IM-sequence length checkpoint from (Nguyen et al., 2023). For M2-BERT-base, Hyena-s-4K, and HyenaDNA, we additionally use fast depthwise convolution kernel for short kernels. For M2- Be--BETT-basy base, Hyena-s-4K, and HyenaDNA, we report results benchmarked on one H100- -X.. For the others, we report performance on one A A100-40GB.
10M model from Monarch Mixer (Fu
D.6 COMPARISON TO TRANSFORMERS
For the comparison against Transformers Table 5, we use the official implementations with the FlashAttention-Vy -2 release (Dao, 2023). We use Hyena model, and match the number of layers, hidden dimension, and expansion factor to the 2.7B Transformer model. To compute the FLOP usage, we take the formula:
2 tokens num parameters
for the parametric FLOPs. For the non-parameter FL.OPs, we add the raw FLOP count from our cost model in Equation 2 (without the adjustment for speed of tensor core FLOPs).
D.7 PARTIAL CONVOLUTIONS FOR HYENA
For the measurement of memory footprint reduction in Table 6, use the same Hyena-s model as in Tables and 4, except we cut the filter short. This lets us offload parts of the input, which reduces the memory footprint.
D.8 EXTENDING HYENADNA-IM
In Table 7, we use a sliding window approach to extend the HyenaDNNA -1 1M and HyenaDNA-4550K models to longer sequences. This mimics training a 4M-sequence HyenaDNA with short filter.
D.9 FREQUENCY-SPEN SPARSE CONVOLUTIONS
To evaluate frequency rreaueecy-pparse sparse convolutions, take the pretrained HyenaDNA- 1M model, and sparsify ky using the strategy described in Appendix B.4. We then run standard validation using the validation set from (Nguyen al., 2023).
28

---

Published as a conference paper at ICLR 2024
Table 21: Measured Constants for Cost Model for A A100-40GB.
Constant
A100-40GB
IB35
08
TB5s
TM
234 TFLOPs
7G
17.6 TFLOPs
D.10 EMPIRICAL GPU PROFILING
Table gives empirically-menn meeauurre GPU stats for A100-40GB, which we used to generate Figure The statistics are specialized to the Monarch decomposition workload. To measure the achievable tensor core FLOPs, we measured the utilization real fpl6 matrix multiply. To measure achievable general arithmetic FLOPs, we measured the utilization of continuously applying Twiddle factors. To measure the achievable HBM bandwidth, we measured the speed of torch. clone a tensor. To measure the achievable SRAM bandwidth, we measured the slow down from writing intermediate results to SRAM between matrix multiply instructions.
29


Reviews:
[{"review_id": "4Fe5WNpdjU", "reviewer": "Reviewer_pq6T", "summary": "This paper proposed to improve the Fast Fourier Transform in the long convolution operation by better utilizing hardware advances. Concretely, it first decoupled the matrix to smaller components such that they can be computed via matrix multiply units (Tensor Cores) in hardwares. Second, through monarch decomposition, it allows the proposed algorithm to process much longer sequences under the constrains of shared memory of GPUs.\n\nExperiments on Hyena and M2-Bert demonstrate that the proposed algorithm achieved significant efficiency, while does not hurt model performance.", "questions": "One question is about the precision of the proposed algorithm: since all the intermediate computations are performed using bf16/fp16, I was wondering how precise they are when comparing with the vanilla FFT implementation under fp32.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "1. This paper aims to solve an import problem in long-sequence convolution computation: poor hardware utilization for FFT.\n\n2. The proposed algorithm achieved significant efficiency, while maintaining model accuracy/performance.", "weaknesses": "NA.", "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": "DeYQlcn21x", "reviewer": "Reviewer_Bh63", "summary": "In this paper, the authors propose an efficient method to compute the convolution of long sequences by exploiting new computational features available on modern GPUs, tensor cores. The motivation for this study comes from the long sequences that are typically encountered in language and time-series models that require large filters, in stark contrast with the usually small filters utilized in convolutional vision models. The method is based on a matrix interpretation of the FFT algorithm, specifically a p-Monarch formulation, that decomposes the FFT operations into a small set of multiplications that are performed using the much higher arithmetical intensity afforded by tensor cores on modern A100 and H100 GPUs. The authors discuss several algorithmic details required to achieve high performance using this FFT expression, such as parallelizing over the sequence, instead of a direct distribution of work over the batch and hidden dimension and exploiting several properties related to the nature of the FFTs of interest to large sequence models, such as folding a K element real-to-real transform into a K/2 complex transform. With these observations, the authors were able to significantly extend the applicability of the proposed method to sequences up to 32K 16-bit entries processed within a single threadblock. Performance studies on long sequences demonstrate the performance improvements achieved compared to the baseline FFT implementation available in PyTorch and a fused version implemented using the cuFFTdx library.", "questions": "- Are the arrows in Figure 3 in the correct positions? The text references tradeoff or crossover points between the different order-p decompositions being of interest but 2 of the arrows seem to reference downward slopes. Maybe I'm interpreting either the text or the graph incorrectly.\n- Does the cufftDx variation use also fold the inputs to perform an order N/2 complex transform instead of an order N real-to-real transform?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Generally well-written with an extensive presentation of the algorithmic and system-related details provided in the appendices.\n- Highlights a key difference between traditional convolution filter sizes present in vision models versus larger models and utilizes a decomposition that takes into account modern architectural features available on GPUs, tensor cores specifically.\n- The proposed strategy to process the inputs over sequences effectively maximizes the number of elements that may be processed in a single thread block and maximally utilizes the tensor cores with more work per block.\n- By fusing multiple operations into a single kernel the authors increase the arithmetic intensity of the conv kernel to move it away from the memory boundness of a single FFT call to a more compute-rich fused counterpart.", "weaknesses": "Major:\n- I would appreciate a clearer distinction between the order-p Monarch FFT and the classic (Bailey) 4-step FFT. The 4-step FFT may also be computed using the Fourier matrices in $\\mathcal{O}(n^2)$ work but, as noted in the text, this expense is reduced substantially by realizing a DFT will perform the same action in $\\mathcal{O}(n \\log n)$ work. It's not clear from the description how the reader to reconcile these 2 competing ideas regarding the reduction in computational work and improvement in overall performance. I suppose the idea is that although the algorithmic expense of the matrix formulation is higher the small sizes of the blocked inputs coupled with the increased throughput of the tensor core units vs the general arithmetic path makes the overall algorithm performance profitable, is that the correct way to think about this issue?\n\nMinor:\n- HBM in section 2.2 introduced as global memory instead of high-bandwidth memory\n- I'm not convinced Figure 2 adds any meaningful insight into the differences between the different broadcasting strategies. Parallelization over the batch and hidden dimensions seems clear but the alternate strategy to parallelize over sequences is less informative.\n- $\\sigma_H$ and $\\sigma_S$ are defined in section 3.3 but never referenced in the cost model or definition of $w(i)$.\n- The components of the cost model could use a bit more explanation to ensure the reader is aware of the origin of each component and its relationship to the algorithmic definition.\n- Is the precision of the datatypes ever mentioned in the main text? I see it is referenced in Appendix D as a 16-bit type.", "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": "cK3q1jlpgp", "reviewer": "Reviewer_8BY8", "summary": "The paper proposes FlashFFTConv, a new algorithm for computing FFT on GPUs. This algorithm is more efficient than existing methods and can be used to accelerate a variety of machine learning tasks, especially for the long-sequence tasks. The author proposes approximation algorithms by leveraging the sparsity.", "questions": "1. What are the limitations and extensions of the method?\n2. What is the potential negative impact of the method? Can we always obtain better performance on FlashFFTConv over cuFFT?", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well organized and easy to follow.\n2. The implementation is in the supplementary material. The appendices provide many details.\n3. The method is solid and the experiments are convincing.", "weaknesses": "### Major issues\n1. My understanding is that the main algorithm is arithmetic equivalent to the standard FFT. In other words, the proposed methods can be seen as an implementation of the FFT without any approximations. The method in Section 3.3 is an approximation algorithm. Is it correct?\n2. FlashAttention is a great success since the attention is a \"new\" operator and lack efficient implementation. However, FFT is a standard operator with long history. Specifically, the Monarch FFT Decomposition was developed in the last century. Considering that this algorithm is not new and GPU engineers can handle the memory hierarchy and other hardware specifications properly, why this method was not developed previously?\nI do not know the detailed algorithms in the cuFFT library, which is assumed to be highly-optimized. What is the reason of better performance of FlashFFTConv over cuFFT, (1) the better algorithm, i.e., Monarch FFT Decomposition, or (2) a better implementation considering the GPU architecture?\n3. In Table 3, FlashFFTConv outperforms torch.fft by up to 8.7×, while the speedup is about 2x without the domain-specific optimizations. Does it mean the major speedup comes from the domain-specific optimizations instead of the FlashFFTConv algorithm? Could the authors conduct this ablation study (with and without the domain-specific optimizations) in other experiments?\n4. When analyzing Table 4, the authors claim that \"speedup varies by the size of the models and the relative amount of time spent computing the convolution compared to other parts of the models\". Please provide the quantitative results, e.g., the relative amount of time spent computing the convolution.\n\n### Minor issues\n1. What about the results on small and medium sequence tasks?\n2. Other than the machine learning applications, FFT is widely used in many fields. It is better to expand the scope of the paper. What about the results in signal processing?", "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)