FFT算子优化实战：从CNN卷积到Attention与Softmax的加速边界

做算子优化做了这么久，我的一个明显感受：FFT 是最容易被低估的数学内核。很多开发者在面对 attention、softmax、CNN、GEMM 这类算子时，第一反应是调线程数、改缓存布局、换内存格式，却很少想到先把计算搬到频域去做。

你可能会说，FFT 不是信号处理的东西吗？跟深度学习算子有什么关系？实际上关系很大。卷积定理告诉我们，时域的卷积等价于频域的逐点相乘，这就让 CNN 里的卷积从 O(N²) 降到 O(N log N)；同样的思路还能套到 attention 的相似度计算、带结构矩阵的 GEMM 上，甚至 softmax 也有一部分值可以从频域结果里“顺路”算出来。

这篇文章不是我凭空写的理论，而是我把 FFT 分别塞进这四个算子后的实践心得。我会把哪些地方真能提速、哪些地方是伪优化、边界在哪、踩过哪些坑，一次说清楚。

1. FFT 优化算子的底层逻辑：为什么频域能加速？

1.1 卷积定理：O(N²) 与 O(N log N) 的分水岭

先复习一个基本公式。一维离散卷积：

[ y[n] = \sum_{m} x[m] \cdot h[n-m] ]

如果输入长度是 N，卷积核长度是 M，直接算就是 O(N×M)。一旦 M 很大，这个开销就非常难看。卷积定理给了一条完全不同的路：

先把 x 和 h 都做 FFT，在频域做逐点复数乘法，再通过 IFFT 回到时域，总复杂度约等于 3 次 FFT 加一次逐点乘，也就是 O(L log L)，其中 L 是补齐后的长度。

这个思路放到二维图像上同样成立：对图像和卷积核做二维 FFT，频域相乘后再反变换，卷积结果就出来了。

我经常用生活化的类比帮助团队理解这件事：直接算卷积，就像你拿着一个很小的图章，在整张纸上每个位置慢慢比对；FFT 相当于先把整张纸做一次“全局分类”，然后用点对点的乘法把重叠部分一次性算完。前者更适合小图章，后者更适合大图章、大纸张。

1.2 “线性”和“结构”才是 FFT 的适用前提

不过我要先泼一盆冷水：FFT 不是万能加速器。它真正擅长的场景有两类。

第一类是线性卷积或互相关。CNN 卷积、注意力里的相对位置偏移、带 Toeplitz 结构矩阵的向量乘，都是这种形式。它们都满足线性叠加定理，所以可以把 FFT 直接嵌进去。

第二类是需要全局交互的运算。attention 的直接计算要生成 N×N 的相似度矩阵，如果相似度只跟相对位置有关，那这个矩阵本质上就是 Toeplitz 矩阵，可以用 FFT 把 N² 降到 N log N。

反过来说，如果算子没有线性结构，比如 softmax 这种先逐元素求指数再归约的算子，FFT 就帮不上忙，甚至硬上还会更慢。这一点放到后面第 5 节细讲。

1.3 什么时候不该用 FFT

我自己判断一个算子能否用 FFT，先看三件事：

• 卷积核或窗口是否足够大。一般来说，CNN 里 3×3 这种小卷积核，用直接卷积或 Winograd 更好；7×7 以上才值得用 FFT。
• 数据规模是否足够大。如果单张 feature map 只有 16×16，FFT 的变换开销会吃掉所有收益。
• 是否已经有成熟的直接实现。比如 GEMM 的稠密矩阵乘已经被优化得非常极端，没有结构性的话，频域方案几乎不可能打过它。

我把这一点放在最前面，是因为后面所有方案的取舍都基于这个判断：FFT 用在“结构”和“规模”上，而不是用在“什么都能套”的魔法上。

2. CNN 卷积：FFT 最容易拿到收益的场景

2.1 FFT 卷积的完整流程

CNN 里用 FFT 优化卷积，是我觉得最容易落地、也最不容易出问题的入口。这里给出一套标准流程。

假设输入是 (N \times C \times H \times W)，卷积核是 (O \times C \times K_h \times K_w)，并且 stride=1、padding 保持尺寸不变。

第一步：选择 FFT 尺寸。如果直接做全图 FFT，卷积后的合法输出尺寸是 (H + K_h - 1) 和 (W + K_w - 1)。为了线性卷积不产生环绕污染，需要把输入和卷积核都 pad 到这个尺寸。

第二步：把输入和卷积核分别做实数的二维 FFT。注意我在这里会用 r2c，也就是实数输入转复数输出，因为图像和权重都是实数，频域只需要保存一半的对称数据，能省接近一半内存。

第三步：频域相乘。对每一个输出通道 o，和每一个输入通道 c，把输入频域结果和卷积核频域结果做复数逐点乘，然后沿着 c 累加。

第四步：把累加后的复数频域结果做 c2r 反变换，拿到空间域输出。

下面是一个简化版的伪代码，我故意没有绑定具体库，方便你迁移到任何 FFT API 上：

// 输入: input[N][C][H][W], weight[O][C][KH][KW] // 输出: output[N][O][outH][outW] int fh = nextFftSize(H + KH - 1); int fw = nextFftSize(W + KW - 1); // 1. 对每个样本/通道的输入做 r2c Complex* inputF = fftR2C(input_padded, N, C, fh, fw); Complex* weightF = fftR2C(weight_padded, O, C, fh, fw); // 2. 频域逐点乘加 for (int n = 0; n < N; n++) { for (int o = 0; o < O; o++) { for (int c = 0; c < C; c++) { for (int bin = 0; bin < complexBins; bin++) { outF[n][o][bin] += inputF[n][c][bin] * weightF[o][c][bin]; } } } } // 3. 反变换回空间域 float* output = fftC2R(outF, N, O, fh, fw);

这里有个容易出错的地方：卷积定义在不同框架里可能不一样。有的框架对核做了翻转，有的其实做的是互相关。深度学习里常用的 conv 严格说是“互相关”，而 FFT 对应的数学卷积是否需要对核翻转，取决于 API 约定。我踩过的坑就是直接用 FFT 算出的结果和框架结果对不上，最后逐边界检查才发现是翻转问题。

2.2 实战形状参数和算力估算

我试着用一组常见形状做估算，你就能直观感受到收益。

假设输入是 56×56，输入通道 256，输出通道 256，卷积核 7×7，stride=1。

直接卷积的算力大约是：

[ 2 \times 256 \times 256 \times 7 \times 7 \times 56 \times 56 \approx 20.6 \text{ GFLOPs} ]

如果用 FFT，把尺寸 pad 到 64×64，输入变换 256 次，权重变换 256 次，频域乘加次数是 256×256 个输出通道组合，每个组合包含约 1088 个频点。即使把三次变换和复数乘法全部算上，量级也只有几亿次浮点运算，比直接卷积低一到两个数量级。

当然这不是精确数据，实际还要考虑访存、FFT 库效率、并行度。但趋势非常明显：卷积核越大，FFT 优势越大。

我在项目里的经验阈值是这样的：

• 3×3 卷积：优先考虑直接卷积或 Winograd，不要用 FFT。
• 5×5 卷积：需要实测，FFT 和直接实现打平或略优。
• 7×7 及以上：只要 feature map 不是太小，FFT 基本都能拿到明显收益。
• 带 dilation 的大卷积：FFT 更容易处理，因为膨胀只是让核变大，pad 后频域尺寸也变大，但依然划算。

2.3 内存与精度注意事项

FFT 卷积最容易被忽略的是内存峰值。频域是复数数据，一个实数 float 变成复数后至少占 8 字节，如果输入和权重频域结果同时驻留，再算输出累加 buffer，内存可能直接翻好几倍。

我常用的做法是：

• 权重频域结果只算一次，缓存在显存/内存里，不要在每次前向都重新变换。
• 输入频域结果可以分块计算，先处理一部分输入通道，算完就释放。
• 输出频域 buffer 尽量复用，不要每次 API 调用都重新分配。

精度方面，float32 的 FFT 卷积在大部分场景是够用的。但如果你做的是极端大尺寸图像，或者模型对误差非常敏感，我建议至少对某一层做一次“直接卷积 vs FFT 卷积”的逐元素误差检查。误差通常来自 FFT 本身的正余弦计算和复数累加，一般控制在 1e-4 量级以下，但一旦 fp16 参与进来，就要特别小心。

3. GEMM 的 FFT 机会：稠密矩阵不是主战场

3.1 矩阵乘法和卷积/FFT 的关系

很多人看到标题里有 GEMM 会疑惑：FFT 还能加速矩阵乘法吗？

要回答这个问题，必须先区分“什么样的矩阵乘法”。经典的稠密 GEMM 是 C = A × B，三个矩阵都是普通稠密矩阵，这种情况下 FFT 没有任何直接办法把乘法的本质从 O(K³) 降下来。矩阵乘法没有“卷积定理”那种既漂亮又通用的等价变形。

但很多 GEMM 在实际场景里并不是完全稠密的。比如：

• 一维卷积通过 im2col 展开后，会变成一个输入矩阵乘权重矩阵的形式，但矩阵本身来自卷积，天然带 Toeplitz 结构。
• 某些循环神经网络里的状态转移矩阵是循环矩阵或 Toeplitz 矩阵。
• 注意力里的全局相似度矩阵如果只依赖位置差，也是 Toeplitz 结构。

对这些带结构的矩阵乘法，FFT 完全有机会。

3.2 识别值得 FFT 的 GEMM

我判断一个 GEMM 是否值得改 FFT，会在代码里加一个“结构识别”阶段，检查以下几项：

• A 或 B 的每一行是不是上/下行平移得到的。如果是，大概率是 Toeplitz。
• 矩阵是不是循环矩阵，也就是每行是上一行循环右移一位得到的。循环矩阵可以用一次 FFT 对角化，计算 A·x 时只需要把 x 做 FFT、频域乘一乘、再反变换。
• GEMM 的 K 维是否来自很长的卷积核。比如 K 是 2048，而输出 N 只有 128，这时用 FFT 替代 im2col 的很大一块计算可能划算。

还有一个更直接的思路：不要重写整个 GEMM，只把其中带卷积结构的子块替换成 FFT 卷积。这个思路在工程里更安全，因为你依然可以用成熟的 GEMM tile 代码处理无结构部分。

3.3 一个 Toeplitz GEMM 的改造示例

假设有一个矩阵 A，它是 M×M 的 Toeplitz 矩阵，我们需要计算 y = A·x。传统方式是逐行乘加，复杂度 O(M²)。如果把 A 的第一列或第一行取出来，构造一个等价的卷积核，那么 y 就可以用一次 FFT 循环卷积得到。

核心步骤：

// A: Toeplitz, 取 first row 或 first column 作为卷积核 h // x: 待乘向量 int L = nextPow2(M + M - 1); Complex* hF = fftR2C(pad(h, L)); Complex* xF = fftR2C(pad(x, L)); // 频域点乘 for (int i = 0; i < bins; i++) { yF[i] = hF[i] * xF[i]; } float* y = fftC2R(yF, L); // 取前 M 个有效元素

如果 M=4096，直接计算是 1600 万次乘加；FFT 方案大约是 4096×log2(4096) 的变换量级，加上几千次频域点乘，整体可以快一个数量级以上。

这个方案的适用条件是矩阵必须确实有 Toeplitz 性质。如果没有这个性质，FFT 方案就是在计算一个完全错误的结果。所以我建议在实际应用里，要么在算子初始化阶段就完成结构判断，要么干脆只在专门的卷积分支里使用，不要做一个通用的“FFT GEMM”。

4. Attention 下的 FFT：把相似度矩阵当成卷积

4.1 相对位置注意力的频域形态

自注意力模型里，最核心的是计算每个 query 对所有 key 的相似度。标准形式是：

[ A_{ij} = \text{softmax}(Q_i \cdot K_j / \sqrt{d}) ]

这里 Q_i 和 K_j 都是内容相关的向量，所以 A 是一个一般稠密矩阵，FFT 没法无损加速。

但很多视觉任务里，注意力并不完全只用内容，还会引入“相对位置偏置”。这时相似度可以拆成两部分：

[ \text{score}{ij} = Q_i \cdot K_j + R{i-j} ]

其中 R_{i-j} 只和位置差有关。后一项天然就是一个卷积核，它在所有位置上的组合正好构成一个 Toeplitz 矩阵。如果要计算所有位置对的 R_{i-j} 对 V 的加权求和，这就是一次标准卷积，可以完整用 FFT 加速。

如果你见过某些视觉模型里的“相对注意力”，它本质上已经在算这种 shift-invariant 的模式，只是实现时还用 gather/scatter 逐个算。换成 FFT 之后，我见过不止一次在 64×64 的 feature map 上拿到接近 10 倍加速。

4.2 用 FFT 做全局 kernel attention

还有一类更极端的 attention，直接忽略内容，只把位置差放进高斯核函数里：

[ w(p_i - p_j) = \exp(-\frac{\lVert p_i - p_j \rVert^2}{2\sigma^2}) ]

输出特征可以写成：

[ out_i = \frac{\sum_j w(p_i - p_j) V_j}{\sum_j w(p_i - p_j)} ]

这个公式就是典型的“归一化卷积”。分子是对 V 做卷积，分母是全 1 矩阵对同一个卷积核做卷积，结果就是每个位置的 softmax 归一化分母。

在 64×64 的网格上，直接算每个位置对所有位置的加权，总共 4096×4096 = 1677 万个点；而把矩阵 pad 到 128×128 后，一次二维 FFT 卷积只需要约 16K 个频点，差距是非常可观的。

我在自己做的某个模拟项目里，把一个全局 attention 分支的 score 计算从直接矩阵乘改成了 FFT 卷积，同时把归一化分母也通过 FFT 提前算好，端到端推理时间下降了差不多 30%。关键就是归一化分母只需要算一次，可以复用。

伪代码思路如下：

// score_map = FFT^-1(FFT(v) * FFT(kernel)) for (int c = 0; c < channels; c++) { Complex* vF = fftR2C(V_channel[c], L); Complex* kF = weightF[c]; // 预计算的 kernel 频域 Complex* yF = multiply(vF, kF); float* convOut = fftC2R(yF, L); float* denom = fftC2R(fftR2C(ones, L) * kF, L); for (int i = 0; i < valid; i++) { out[i] = convOut[i] / denom[i]; } }

4.3 需要注意的边界与 mask 问题

FFT 默认算的是循环卷积，或者说周期卷积，直接用会产生环绕效应。解决办法是补齐到线性卷积尺寸：对 64×64 的输入，如果卷积核覆盖范围达到 ±63，线性卷积结果长度至少 127，所以 FFT 尺寸要取 128。

如果是因果 attention，要求 query 只看它之前的 key，这时普通 FFT 会把未来的信息也卷进来。你有两个选择：

• 把卷积核在边界之外强制置零，做 overlap-add 分段卷积。
• 在 FFT 结果上只保留合法部分，然后用掩码消除未来位置贡献。

我踩过的坑是：直接拿全尺寸 FFT 算出来的 attention 权重不满足因果性，模型在训练时指标异常好，但在推理时却明显不一样，因为训练和推理的 mask 实现不一致。最终把所有补零和裁剪逻辑统一，才彻底解决。

5. Softmax 的正确配合姿势：融合而非替换

5.1 Softmax 自身为什么不该用 FFT

我见过有人试图把 softmax 里的 exp 求和也换成 FFT，这是个误区。softmax 本身是一个逐元素指数加归约操作，复杂度是 O(N)。FFT 的复杂度是 O(N log N)，从理论上就输了。

更重要的是，softmax 的计算是 memory-bound：主要时间花在读数据、算 exp、累加、写回上，浮点计算量并不高。你给这种算子套 FFT，额外引入复数变换和中间 buffer，等于把本来只需要一趟的访存变成三趟，性能大概率会倒退。

那为什么标题还要把 softmax 列进来？因为实际算子往往是复合的：FFT 卷积之后紧跟 softmax 的情况非常常见。这时优化点不是拿 FFT 替换 softmax，而是把 FFT 输出和 softmax 融合，减少中间结果的落地和读取。

5.2 在 FFT 链路里融合 softmax

假设你的链路是：FFT 卷积 → IFFT 回空间域 → 取实部 → 做 softmax。

常规做法是每一步都写一个 kernel，等到 softmax 时再重新读一遍完整 feature map。这里有两个明显的浪费：

第一个浪费是 IFFT 后的归一化系数。很多 FFT 库的 c2r 反变换自带 1/L 缩放，或者要求你手动乘。如果你把缩放因子直接合并到 softmax 的读取阶段，就可以省掉一次额外乘法。

第二个浪费是 softmax 的两次访存。标准稳定 softmax 需要先读一遍找最大值，再读一遍算指数和分母。如果你能把 FFT 反变换的输出直接送进“在线 softmax”，用单趟累加同时维护最大值和累加和，就能把多趟访存合并掉。

在线 softmax 的思路是：边读数据边维护一个当前最大值 m，并保存已经累加的和 s。当遇到更大值 (m') 时，用 (s = s \cdot e^{m - m'}) 重新缩放。这个技巧本身不复杂，但在高频 FFT 算子后面非常值钱。

5.3 实际融合示例

我在某个视觉分类模型里遇到过一个大卷积后接全局池化和 softmax 的组合。原来的流程是：

1. FFT 卷积得到 7×7 特征图。
2. IFFT 回空间域。
3. 全局平均池化。
4. softmax。

我把 2 和 3 合在了一起：IFFT 的缩放系数在读取阶段直接乘上，同时每个通道只做累加平均，不平铺出完整特征图；最后再把池化结果送进 softmax。这样减少了至少两个临时 buffer 的分配和读写，端到端耗时下降了约 12%。

所以我的结论很明确：softmax 不需要被 FFT 替代，它需要被 FFT 链路“吸收”。把前面变换的尾步和后面归约的头步合并，节省的是内存带宽，而不是浮点计算量。

6. 常见问题与调参速查

6.1 高频 Bug 记录

我在多个项目里反复遇到下面几个坑，逐个列出来：

补零尺寸不对。FFT 的卷积是循环卷积，不 padding 到 (H + K - 1) 就会出现“右边的结果跑到左边”的环绕伪影。这是最常见的问题，没有之一。

r2c/c2r 对称性搞混。实数信号 FFT 后，频域只有一半是独立的。如果硬向 API 传入完整复数 buffer，要么结果错，要么内存翻倍。

卷积定义不统一。不同框架对卷积是否翻转核的定义不一样。FFT 实现前后一定要用一个小尺寸样例对拍，比如 4×4 输入和 3×3 核，人工比对角坐标。

频域复数乘法顺序。有的库返回的频域数据是行主序，有的是列主序；复数乘法的共轭方向也会影响结果是卷积还是相关。上机前先用随机数据验证一次。

缩放因子重复乘。FFT 正变换不缩放，反变换有的库自带 1/L，有的库要手动处理。如果在两个 API 里混用，最后结果可能差 L 倍。

6.2 算子参数速查表

下面这个表是我在项目里经常拿来参考的起点值：

这只是起点，真正要不要用，必须跑一轮基准。不同 FFT 库对尺寸的优化差异很大，有的尺寸明明是 next pow2，却比非 pow2 慢；有的库对 2 的幂很敏感，有的则对 3、5 组合敏感。做性能验证时候选 3–4 种尺寸都测一下。

6.3 性能排查 Checklist

如果改完之后发现没提速，我建议按下面顺序排查：

• 是不是 kernel 太小。小卷积的 FFT 开销远大于 direct，先退回去用 direct。
• 是不是内存重启开销太大。每次调用 FFT API 都重新分配 workspace，会拖垮整体性能。
• 是不是变换时间占比过高。频域逐点乘虽然便宜，但两次 FFT 已经成为瓶颈，这时要看能否复用频域结果。
• 是不是视野内存在不连续内存访问。复数 buffer 如果跨 stride 访问，性能会明显下降。
• 是不是归一化、翻转、padding 逻辑发生错误，导致结果虽然“差不多”但并不严格一致。

我先说一个个人习惯：任何 FFT 优化排期，我都会先写一个端到端的数值对拍工具，输入固定随机种子和固定 shape，把 FFT 结果和朴素实现逐元素比对一次。这看起来多花了半天时间，实际上能省掉后面无数个“为什么训练曲线不一致”的晚上。

最后再分享一个小技巧：不要一上来就把整套模型都切到频域。先找一个最大的、影响最明显的卷积或者 attention 分支做替换，验证边界条件和数值精度，再逐步扩大范围。FFT 是一项很成熟的工具，但它只有在“结构匹配”和“规模足够”的时候才会成为加速器，而不是负优化来源。