Paper: FlashFFTConv: Efficient Convolutions for Long Sequences with Tensor Cores (ICLR-2024)

Reviewer comment (Theory):
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?

Author response:
**Approximation Algorithms**
We clarify that there are actually **three algorithmic contributions** in our paper. The main contribution is the FlashFFTConv algorithm, which is an **exact** FFT convolution algorithm.

But the structure of the FlashFFTConv algorithm allows immediate applications to sparsity – simply skipping portions of the matrix multiply operations yields sparse convolution algorithms. Partial convolutions and frequency-sparse convolutions are two approximate convolution algorithms that are enabled by the computational pattern of FlashFFTConv.

Reviewer comment (Novelty):
FlashAttention is a great success since the attention is a "new" operator and lacks efficient implementation. However, FFT is a standard operator with a 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 was this method not developed previously?

Author response:
**Why is this not in cuFFT?**
In short – a single FFT on its own is almost always memory-bound, so the most efficient FFT implementation does not need to use tensor cores or other optimization. In other words, the older FFT algorithms suffice. These are what cuFFT uses. Please see the common response for a more detailed discussion on the history of FFT algorithms, and why the FFT convolution presents specific bottlenecks. 

When fusing multiple FFT operations together, the entire operation becomes compute-bound, so different algorithms are necessary – which is why we need to translate the operation to use tensor cores.

Finally, we note that the use case of long convolutions has started to show promise in ML applications only relatively recently. For example, S4 was only published at ICLR last year, and gated convolutional models that come close to attention performance in language modeling have only started to emerge this year. So compared to FlashAttention, which took 5 years between the introduction of attention and the first IO-aware algorithm, FlashFFTConv is moving relatively quickly!

Reviewer comment (Evaluation):
I do not know the detailed algorithms in the cuFFT library, which is assumed to be highly optimized. What is the reason for the better performance of FlashFFTConv over cuFFT: (1) the better algorithm, i.e., Monarch FFT Decomposition, or (2) a better implementation considering the GPU architecture?

Author response:
**Why is this not in cuFFT?**
In short – a single FFT on its own is almost always memory-bound, so the most efficient FFT implementation does not need to use tensor cores or other optimization. In other words, the older FFT algorithms suffice. These are what cuFFT uses. Please see the common response for a more detailed discussion on the history of FFT algorithms, and why the FFT convolution presents specific bottlenecks. 

When fusing multiple FFT operations together, the entire operation becomes compute-bound, so different algorithms are necessary – which is why we need to translate the operation to use tensor cores.

Finally, we note that the use case of long convolutions has started to show promise in ML applications only relatively recently. For example, S4 was only published at ICLR last year, and gated convolutional models that come close to attention performance in language modeling have only started to emerge this year. So compared to FlashAttention, which took 5 years between the introduction of attention and the first IO-aware algorithm, FlashFFTConv is moving relatively quickly!

Reviewer comment (Experiments):
In Table 3, FlashFFTConv outperforms torch.fft by up to 8.7×, while the speedup is about 2× 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?

Author response:
**More Detailed Breakdown of Speedup**
Thank you for the feedback to conduct a more detailed ablation study of the speedups. We have re-organized our convolution benchmarks to several distinct use cases:
* A pure convolution, where the input and FFT size are the same length
* A convolution with padding, where the input is half the length of the FFT size (i.e., causal convolution)
* Gated convolutions, where there are gating operations before and after the convolution to fuse

We have chosen to split the benchmarks into these distinct use cases, instead of a single table with “domain-specific optimizations.” The full results are given in Appendices C.1 and C.2 of the updated draft.

Reviewer comment (Evaluation):
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.

Author response:
**Percentage of Time Each Model Spends Computing the Convolution**
We provide the percentage of time each model spends computing the convolution when implemented in PyTorch, and the relative speedup from using FlashFFTConv:

|      **Model** | **% Time in Convolution** | **Speedup** |
|---------------:|:-------------------------:|:-----------:|
|    **M2-BERT** |           48.5%           |     1.9x    |
|      **Hyena** |           38.4%           |     1.7x    |
| **Long Convs** |           64.3%           |     2.4x    |
|    **SaShiMi** |           22.1%           |     1.3x    |
|   **HyenaDNA** |           81.8%           |     4.4x    |

As you can see, SaShiMi spends the least amount of time computing the convolution, and thus has the least amount of speedup. HyenaDNA spends the most amount of time in the convolution, and has the most speedup. We have added this table to Appendix C.

Reviewer comment (Experiments):
What about the results on small and medium sequence tasks?

Author response:
**Results on Small and Medium Sequence Tasks**
M2-BERT is a small-sequence task (sequence length 128, benchmarked in Table 4 of the original submission). We also provide benchmarks of Hyena-2.7B at different sequence lengths (2K-16K), and Hyena-s-4K, which are both medium-sequence tasks.

Reviewer comment (Experiments):
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?

Author response:
**Applications Outside Machine Learning**
The FFT is very widely used outside of machine learning, so we look forward to applications outside of machine learning as exciting future work for this project. We have chosen to focus on machine learning applications for this paper, since ICLR is a machine learning conference. However, we note that S4/SaShiMi and Hyena, which we benchmark in the paper, are built on signal processing primitives. We also evaluate on HyenaDNA, a scientific application, and we are sure there are many others that may be relevant.

Reviewer comment (Evaluation):
What is the potential negative impact of the method? Can we always obtain better performance on FlashFFTConv over cuFFT?

Author response:
**Potential Negative Impact**
For FFT convolutions, FlashFFTConv is faster across all sequence lengths. However, FlashFFTConv does currently not support fp32 precision, which could be necessary for some applications with high-precision requirements. Some FFT applications in science may require double precision, which FlashFFTConv currently does not support.

Reviewer comment (Evaluation):
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?

Author response:
**Differences between Monarch, Bailey, and 4-step FFT.** Mathematically, all these algorithms are equivalent to each other. For a more in-depth discussion, Van Loan [1] is a great resource on how different FFT algorithms are almost always mathematically equivalent, but use different decomposition strategies for computational reasons.

Historically, Bailey’s FFT algorithm has focused on the **data movement** aspects of the FFT computation – it was designed for settings where the entire input cannot fit into the local memory of a machine. In this way, it is similar to a streaming multiprocessor on a GPU; a single streaming multiprocessor has a limited amount of SRAM available to store the input, so longer sequences require communication to HBM.

The major distinction in the Monarch decomposition is the emphasis on the **compute units** used to compute the individual FFT operations. In the Monarch decomposition, the individual FFT operations are computed using a dense matrix multiply operation. As the reviewer notes, this incurs higher absolute FLOP cost than an $O(N log N)$ algorithm – but the increased throughput of the tensor core units makes up for it.

In FlashFFTConv, we are careful to balance the sizes of the matrix multiply operations with the tensor core units available; we usually only compute 16x16 matrix multiplies, or 32x32 (the tensor cores are 16x16-sized). Larger matrices incur too much additional FLOP cost to be worth it - this is why the blue line in the cost model in figure 3 rises above the rest.

Reviewer comment (Presentation):
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.

Author response:
**Sequence Length vs. Batch/Head Parallelism.** 
Thank you for the feedback on Figure 2. Our main point is that we can express multiple parallel FFT operations as a single matrix-matrix multiply operation. We have updated Figure 2, and we welcome additional feedback about its clarity.

Reviewer comment (Writing):
$\sigma_H$ and $\sigma_S$ are defined in section 3.3 but never referenced in the cost model or definition of $w(i)$.

Author response:
**Cost Model**
Thank you for the feedback on the clarity. We have added additional exposition around the components of the cost model. $\omega(i)$ is just a helper function that lets us simplify the expression a little by not needing to separate out HBM and SRAM intermediates into different functions.

$\omega(i)$ returns $\sigma_H$ if the intermediates of step $i$ of the decomposition need to be stored in HBM, and $\sigma_S$ if they can be stored in SRAM.

Reviewer comment (Writing):
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.

Author response:
**Cost Model**
Thank you for the feedback on the clarity. We have added additional exposition around the components of the cost model. $\omega(i)$ is just a helper function that lets us simplify the expression a little by not needing to separate out HBM and SRAM intermediates into different functions.

$\omega(i)$ returns $\sigma_H$ if the intermediates of step $i$ of the decomposition need to be stored in HBM, and $\sigma_S$ if they can be stored in SRAM.

Reviewer comment (Presentation):
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.

Author response:
**Arrows in Figure 3**
These arrows refer to the “humps” for $p=3$ and $p=4$ for short sequence lengths – cost is high even though overall FLOP cost is low, since those FLOPs do not fully utilize tensor cores. We have updated the figure in the draft to try to make this more clear, and welcome further feedback.

Reviewer comment (Reproducibility):
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?

Author response:
**cufftDx Variant**
Yes, the cufftDx variant also uses the N/2 trick for the real-to-real transform.

Reviewer comment (Evaluation):
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.

Author response:
We thank the reviewer for the positive feedback. For the question about the precision of the method, please refer to the table in the common response.