FlashFFTConv: Efficient Convolutions for Long Sequences with Tensor Cores

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.

Strengths

1. The paper is well organized and easy to follow. 2. The implementation is in the supplementary material. The appendices provide many details. 3. The method is solid and the experiments are convincing.

Weaknesses

### Major issues 1. 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? 2. 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? I 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? 3. 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? 4. 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. ### Minor issues 1. What about the results on small and medium sequence tasks? 2. 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?

Questions

1. What are the limitations and extensions of the method? 2. What is the potential negative impact of the method? Can we always obtain better performance on FlashFFTConv over cuFFT?

Rating

6: marginally above the acceptance threshold

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

3 good

Presentation

3 good

Contribution

3 good

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.

Strengths

- Generally well-written with an extensive presentation of the algorithmic and system-related details provided in the appendices. - 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. - 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. - 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: - 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? Minor: - HBM in section 2.2 introduced as global memory instead of high-bandwidth memory - 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. - $\sigma_H$ and $\sigma_S$ are defined in section 3.3 but never referenced in the cost model or definition of $w(i)$. - 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. - 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.

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

Rating

8: accept, good paper

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

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. Experiments on Hyena and M2-Bert demonstrate that the proposed algorithm achieved significant efficiency, while does not hurt model performance.

Strengths

1. This paper aims to solve an import problem in long-sequence convolution computation: poor hardware utilization for FFT. 2. The proposed algorithm achieved significant efficiency, while maintaining model accuracy/performance.

Weaknesses

NA.

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.

Rating

8: accept, good paper

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

3 good

Contribution

4 excellent

Common Response (1/2)

We thank all the reviewers for their time and effort reviewing our paper and for their constructive comments, which have made our paper stronger. We appreciate the positive feedback that our work solves an “important problem” (reviewer pq6T) with a “solid” method and “convincing” experiments (reviewer 8BY8). In this common response, we address some common questions that were brought up by multiple reviewers: a summary of why we need a new algorithm for FFT convolution now (even though FFT has been around for decades) and the precision of the inputs & numerical error. We’re also excited to report that FlashFFTConv has already started to see some adoption to train new models since our initial submission – it is being used to train long-context retrieval models, and to scale up hybrid gated convolution-attention architectures that rival the best open-source LLMs in quality. We're excited by this rapid adoption and impact, and **appreciate your help in making this work better and even more impactful!** # Why do we need to optimize the FFT Convolution? A question brought up by multiple reviewers (Bh63, 8BY8) was why we need to be developing new algorithms for the FFT convolution now, when the FFT has existed in different guises for half a century, and when there have been so many algorithms for the FFT in the past. The key is that we have a **new use case, with new hardware and bottlenecks.** In our case, a recent wave of new ML architectures such as S4, Hyena, BiGS, CKConv, M2, and more [1-5] require optimization of the FFT convolution, and the characteristics of ML accelerators such as GPUs introduce new bottlenecks around I/O and utilizing matrix-multiply units. As an aside – developing new FFT algorithms for new bottlenecks has actually been common in the history of the FFT – the “FFT algorithm” is not actually one algorithm, but many different algorithms. These algorithms are all mathematically-equivalent to each other, but they are computationally different depending on the compute and memory characteristics of the systems of the day. For instance, see: * Winograd FFT [6] to minimize multiplication vs. addition (for old hardware platforms where multiplication was more expensive than addition) * Bailey’s FFT [7] was originally designed to distribute extremely long FFTs that did not fit in the local memory of a single machine – and is primarily focused with the memory access pattern. * Hexagonal FFT [8] for data sampled on hexagonal grids FlashFFTConv is another step in this line of new FFT algorithms designed for new hardware bottlenecks. There are two important points: * The convolution requires FFT, pointwise multiply, and inverse FFT. This has a different **memory-compute tradeoff** than a single FFT operation. * Indeed, a **single FFT** does not have this tradeoff, even on GPUs. On GPU, the compute is fast enough (even without tensor cores) that a single FFT is completely memory-bound – so it doesn’t matter if it’s computed with tensor cores or with general arithmetic operations. See below, for a single FFT of length 1K has similar runtime to torch.clone on H100 (and is actually slightly faster, which may indicate slightly better implementation than torch.clone): | **Operation** | **Runtime (ms), B = 64, H = 768** | |----------------:|:---------------------------------:| | **torch.clone** | 0.31 | | **FFT** | 0.27 | | **Fused FFTConv (no tensor cores)** | 0.67 | * When multiple steps of the FFT are fused together, then the memory-compute tradeoff changes, and a **fused FFT convolution is no longer memory-bound** (see the third row above, where a fused FFT convolution takes more time than just reading & writing the output). * That is when it is important to use the tensor cores – compute becomes the bottleneck. That memory-compute tradeoff is why the method is named FlashFFTConv, and not FlashFFT – we specifically focus on the tradeoffs for FFT convolutions on GPU. This is also the **primary difference** between the FlashFFTConv algorithm, and previous statements, like Bailey’s FFT decomposition or the four-step FFT decomposition. The core arithmetic is the same, but previous algorithms emphasize different bottlenecks (i.e., Bailey’s only focuses on the limits of local memory, not the compute units to compute the FFT). Van Loan [9] gives a good overview of frameworks for computing the FFT, and how they are all mathematically the same formulation, but with different compute characteristics, for those interested.

Potential typo in equation (1)

It seems that equation (1) is not convolving over the input--the location i is fixed. It should be $(u \cdot k)[i] = \sum_{j=0}^{N} u[i - j] \cdot k[j]$ ?