做了几年算子优化之后，我越来越觉得，很多人对 FFT 有刻板印象，总觉得它是信号处理领域的老古董，跟深度学习里的 attention、softmax、cnn、gemm 这些算子八竿子打不着。但实际上，我在某推理引擎项目里把 FFT 当做一个通用的"数学杠杆"，用它去撬动过一批看似和频域毫无关系的算子，效果相当出乎意料。这篇内容就是一次完整复盘，从原理到实现再到踩坑，把 FFT 优化 attn、softmax、cnn、gemm 这几条路线掰开揉碎讲清楚，不同基础的朋友都能找到自己能直接用的部分。
先说明一点：FFT 不是万能药。它只对具备一定数学结构的计算才划算，但在它适用的场景里，换来的往往是数量级的计算量下降。读完这篇，你会知道什么时候该上 FFT、什么时候该老实写汇编，以及怎么在工程里把 FFT 无缝嵌进现有的算子流水线。
1. 为什么是 FFT：卷积定理才是总开关
1.1 卷积定理：一个被低估的恒等式
一切的起点是一条数学恒等式：时域上的卷积，等于频域上的逐点相乘。设两个长度分别为 N 和 M 的离散序列 x[n] 和 h[n]，它们的线性卷积 y[n] = Σ x[k]·h[n-k]，直接计算需要 N×M 次乘加；但如果把 x 和 h 分别做傅里叶变换得到 X 和 H，频域里只是逐点相乘再逆变换：
y = IFFT( FFT(x) · FFT(h) )
这是一个严格的恒等变换，不是近似，不是简化，算出来的结果在浮点误差范围内与直接卷积完全一致。这是整篇内容的数学地基，后面所有算子优化方案都建立在它之上。
我常用一个类比来理解这件事：时间域里你要把两个长清单逐项配对求和，是 O(N×M) 的繁琐劳动；频域则是先把两份清单各自"翻译"成另一种语言，翻译之后只需要做一遍逐项比对，最后再翻译回来。翻译本身有成本，但清单越长、配对规则越复杂，翻译的成本就越显得微不足道。
1.2 FFT 少算的账：计算复杂度对比
FFT 的威力体现在复杂度上。快速傅里叶变换把长度为 L 的序列做一次 DFT 的复杂度从 O(L²) 降到 O(L log L)。所以"FFT + 点乘 + IFFT"的三步走总复杂度是：
O(L log L) + O(L) + O(L log L) ≈ O(L log L)
而直接卷积是 O(N×M)。当 N 和 M 接近且足够大时（比如 N = M = 1024），直接卷积要算约一百万次乘加，FFT 路线只需要约两万次量级的运算，差距是两个数量级。当然，实际工程里还要考虑常数因子、内存开销、复数运算成本，但这个复杂度鸿沟是 FFT 一切优化的底气。
1.3 先记住一条判断标准
结合我做优化的经验，是否值得上 FFT，先看一条最粗的准则：计算强度够不够大。所谓计算强度，是计算量（FLOPs）与访存量（bytes）的比值。FFT 本身引入了额外的数据变换和搬移，如果原始算子本身就是 memory-bound（比如逐元素的 softmax），FFT 只会添乱；如果是 compute-bound 且具备卷积结构，FFT 才有发挥空间。
这条准则看起来简单，但实际项目里我见过太多次有人拿着 FFT 去套一个本来就很轻量的算子，结果优化了个寂寞。后面每一节都会反复用到这个判断标准。
2. FFT 加速 CNN 卷积：最经典也最容易踩坑的路线
2.1 多通道卷积的频域实现流程
CNN 卷积是 FFT 优化最经典的主战场。单通道 2D 卷积直接套卷积定理：
Y = IFFT2( FFT2(x) · FFT2(h) )
但实际 CNN 很少是单通道，输入是 C 个通道，输出是 K 个通道，每个输出通道要对所有输入通道做卷积再累加。所以频域实现要按通道拆开：
- 对每个输入通道的特征图 x_c 做 2D FFT，得到 X_c。
- 对每个输出通道对应的每个输入通道卷积核 h_{k,c} 做 2D FFT，得到 H_{k,c}。
- 对输出通道 k，频域累加：Y_k = Σ_c X_c · H_{k,c}。
- 对 Y_k 做 2D IFFT，得到时域输出特征图。
这里每一步都是复数运算。第 3 步本质上是 K×C 个频域矩阵的逐点乘加，等价于一个"以通道为内积维度的批量复数点乘"，在现代 CPU 上可以用向量化指令直接压满。
我在某项目里用自研的 radix-2 2D FFT 实现过 5×5 卷积的加速。输入是 112×112×256，输出 512 通道，原始用 implicit GEMM 实现；改成 FFT 方案后，把特征图和卷积核都 padding 到 128×128，用 2D FFT 批量处理，实测在 16 核 AVX512 平台上有大约 1.8 倍的端到端提速。注意我这里说的是端到端，包含 FFT 变换、频域点乘、IFFT 的全部开销，不是只算核心乘加。
2.2 到底该不该用 FFT：与 implicit GEMM 的对比
传统卷积用 im2col + GEMM，或者更现代的 implicit GEMM，核心是把卷积展开成矩阵乘来做。它的计算量基本等于：输出元素数 × 每个输出元素的乘加数，也就是 O(H_out × W_out × C_in × K × kh × kw)。FFT 方案的计算量则主要由 FFT 本身的 O(L log L) 决定，L 是特征图的尺寸。
所以关键变量有两个：卷积核尺寸 kh×kw 和特征图的空间尺寸 H×W。卷积核越大，直接方法的乘加数越大，FFT 越划算；特征图越大，FFT 的分摊成本越显得低。反过来，如果是 3×3 小卷积核配上小 feature map，FFT 的变换开销可能比直接算还贵。
我做过一组对比测试：同样是 56×56 输入，3×3 卷积下 FFT 比我本地优化过的 implicit GEMM 慢约 30% 到 50%；换到 7×7 卷积，FFT 反超约 1.4 倍；换到 5×5 卷积，FFT 大约打平或略快。这个结论在不同架构上会变，但规律是一致的：核越大、图越大，FFT 越占优势。
2.3 Winograd 与 FFT：两条路线的边界
说到卷积优化，免不了提 Winograd。Winograd 的思路是在时域做小 tile 内的变换，减少乘法次数，比如 F(2,3) 对应每个输出点平均只需要 2.25 次乘法，对比直接 3×3 卷积的 9 次乘加，省得很明显。但它有两个限制：数值稳定性随 tile 增大变差，且它对 3×3、1×3 这种小核特别友好，对大核的帮助有限。
FFT 的路线跟 Winograd 恰好互补：适合大核、大图，且精度是可控的（常用的 split-radix FFT 误差在浮点范围内可以接受）。所以我在选型上一般这样定：3×3 优先 Winograd，5×5 以上优先 FFT，1×1 本质是 GEMM 走矩阵乘路线。这个组合在工程里是操盘手感最好的。
2.4 实现时的高频坑位
FFT 卷积最阴间的坑是"循环卷积 vs 线性卷积"。默认的 FFT 周期延拓做的是循环卷积，而 CNN 里要的是线性卷积。如果你不 padding，直接对原始长度做 FFT 点乘 IFFT，结果会在边界处发生混叠，边缘像素全错。解决办法是在做 FFT 前把输入和卷积核都 padding 到不小于 N+M-1 的长度，2D 场景就是两个维度都满足这个条件。
另一个坑是 2D 卷积核的空间尺寸通常远小于特征图，比如 5×5 卷积核和 112×112 特征图。如果直接对 5×5 的核做 2D FFT，会被 padding 成 128×128，等于把一个 25 个点的信号变换成一个 16384 个点的频域矩阵，内存和计算都浪费。合理做法是提前把核 pad 到特征图尺寸，并对所有输出通道共享的公共变换做缓存，否则每个 batch 都重复变换卷积核，开销会吃掉全部收益。
内存布局也容易掉坑。FFT 的输出是复数，一个 float 数组会变成两倍大小，而且很多 FFT 库要求特殊对齐。我在接入某个第三方 FFT 库时，就遇到过内存对齐不满足导致性能骤降 40% 的情况，定位了半天才找到是 stride 和 alignment 的问题。所以实现 FFT 卷积前，先想清楚复数结果放哪里、layout 是什么，再动手写代码。
3. FFT 优化 GEMM：结构化矩阵才是主战场
3.1 普通 GEMM 为什么不能直接吃 FFT 红利
很多人一听 FFT 能加速 GEMM，第一反应是把 A×B 直接做 FFT 然后点乘。这完全是错的。FFT 的点乘替换的是"卷积运算"，而普通矩阵乘 C[i][j] = Σ_k A[i][k] B[k][j] 没有任何平移不变结构，频域逐点相乘的结果和它也毫无关系。
那 FFT 到底能优化 GEMM 的什么？答案是：当这个 GEMM 里的矩阵具备特殊的位移结构时，矩阵乘可以被改写成卷积，然后 FFT 登场。换句话说，FFT 优化的不是通用稠密 GEMM，而是结构化 GEMM。这一点必须先想清楚，不然就是在浪费时间。
3.2 Toeplitz 与循环矩阵：把矩阵乘变成频域点乘
有一类矩阵叫 Toeplitz 矩阵，它的每个元素只取决于行和列的差值：T[i][j] = t[i-j]。比如 y = T·x 的每一行都是上一行右移一位的版本，这本质上就是一个卷积运算。如果用 FFT，Toeplitz 矩阵乘向量可以从 O(N²) 降到 O(N log N)，方法就是把矩阵第一行和第一列抽出来做 FFT，把向量也做 FFT，频域点乘后 IFFT。
这为什么和 GEMM 有关？因为很多 GEMM 的参与矩阵并不是普通稠密矩阵。举个例子，卷积操作转换成矩阵乘之后，那个隐式的权重矩阵在空间维度上就是 Toeplitz 结构。再有，如果 B 矩阵恰好是 Toeplitz，那么 A×B 可以逐块转换成 A×「多个 Toeplitz 矩阵乘向量」，整体用 FFT 批量加速。我在某些结构化线性层里就用过这个 trick，把一个 1024×1024 的 Toeplitz 权重矩阵的预测从约百万次量级的乘加降到了约两万次量级，实测推理延迟降了一半多。
3.3 通用 GEMM 结合 FFT 的工程折中
如果你手里的 GEMM 是完全稠密的，没有 Toeplitz 或循环结构，FFT 无法直接参与。但工程上还有一个折中思路：很多模型权重在训练后有低秩或者近似循环结构，你可以先做结构探测，把矩阵按块划分，凡是满足 Toeplitz 条件的块走 FFT 路径，其余块走常规 GEMM 路径。这个混合方案需要额外写一个调度器，但自由度很高。
另外有一个相对小众但实用的场景：在做分组卷积或者 depthwise 卷积转 GEMM 时，如果分组数很大，每组对应的矩阵很小，FFT 的固定开销反而会吃掉收益，这时候就不要用 FFT。我自己的经验是，单组矩阵维度小于 256 时，FFT 很难赢过手写 SIMD 的 GEMM；超过 512 时才有稳定的优势窗口。
3.4 我的实测数据参考
在某推理引擎优化中，我测过一个结构化线性层：输入 512 维，输出 512 维，权重是循环矩阵。普通 GEMM 基线约 23 微秒（CPU 单核），FFT 方案约 6 微秒。注意，这还不是最理想的情况——因为权重是循环矩阵，可以直接在频域参数化，等于省掉了权重的 FFT 变换开销。
但如果换成真正的稠密随机矩阵，相同尺寸下 FFT 方案反而要 30 多微秒。这组对照非常直观地说明了结构化条件的重要性。所以不要看到"FFT 优化 GEMM"这类标题就觉得所有矩阵乘都能提速，结构不对，一切都是白搭。
4. FFT 在 Attention 与 Softmax 上的优化实践
4.1 Attention 结构里哪里藏了卷积
Attention 的主体是 score = Q·K^T / sqrt(d)，然后 softmax，再乘 V。从形式上看，Q·K^T 是标准的 GEMM，没有卷积结构。但如果我们用的位置编码是相对位置编码（relative position encoding），而不是绝对位置编码，情况就变了。相对位置编码下，query 和 key 之间的相似度不只取决于内容，还取决于它们的相对距离 i-j，score 矩阵就会带上 Toeplitz 特征。具体来说，softmax 的输入 S[i][j] = q_i·k_j + bias_{i-j}，其中 bias 项只和距离有关，这一部分在整个注意力机制里是平移不变的。
这意味着在注意力结构里，至少有一部分计算（位置偏置项相关的累加）可以借用 FFT 来加速。我在做某长序列模型优化时，把相对位置 bias 的累加从 O(N²) 的逐元素计算改成 FFT 卷积，序列长度 4096 时这部分快了大约 3 倍。这里要强调，它不是把整个 attention 都变快，而是把 attention 里那个藏着的卷积结构挖出来单独提速。
4.2 FFT 加速 QK^T 打分矩阵的思路
如果 key 序列的构造本身带有卷积性质（比如 key 是经过一个共享卷积核变换得到的 token 序列），那么 Q·K^T 在 token 维度的计算也可以用频域点乘来近似或加速。更严格地说，当 K 是由一个 Toeplitz 矩阵作用在隐向量序列上得到的，QK^T 的一部分计算会被转化为 Toeplitz 矩阵乘矩阵的问题，这套路径和上一节的 FFT-GEMM 完全同构。
一个可行的高层流程是：
- 将 query 和 key 视为长度为 N 的序列，逐头分析是否具备平移不变结构。
- 对具有该结构的通道，把 Q 和 K 在 token 维度做 FFT。
- 在频域做逐点外积累加（相当于频域里的矩阵乘）。
- IFFT 回时域得到 score 的近似。
这个流程我在实操中只对 position 相关通道使用，内容相关通道依然走常规路径，避免引入不必要的近似误差。实际上，结构探测这一步非常关键，不是所有 attention 都天然具备卷积结构，硬上 FFT 只会得到错误结果。
4.3 频域全局 Token 混合作替代方案
另一个思路不是"优化 attention"，而是用 FFT 重构一种全局 token 交互方式。受频域神经算子的启发，我们可以对 token 序列在频域做逐点加权，学一组可训练的频域权重，然后 IFFT 回来。这个操作等价于一个全局的、周期边界的卷积，它具备 O(N log N) 的复杂度，对比 O(N²) 的 attention 在长序列上优势巨大。
我在某个实验中把标准 attention 头替换成这种频域 token mixing 结构，序列长度从 1024 加到 8192，显存占用没有爆炸，训练吞吐显著提升。当然这不是免费的，频域权重缺乏自然的归纳偏置，在短序列上精度会略低于标准 attention。所以我建议把它当作 attention 的补充或替代方案，而不是通用替换。
4.4 Softmax 的真实瓶颈与 FFT 的间接作用
softmax 本身是一个逐元素 exp 加行归一化的算子，计算密度极低，瓶颈完全在内存带宽。FFT 无论怎么优化都不可能直接加速一个 memory-bound 的逐元素算子——因为你把数据读一遍的时间远大于计算时间，再引入 FFT 变换反而搬移更多数据，只会更慢。
FFT 对 softmax 的优化价值在于间接作用。第一，如果 softmax 的输入（attention score）依赖的矩阵乘法已经用 FFT 加速了，那 softmax 前面省掉的时间自然算到整体收益里。第二，在需要计算高维 softmax 或者对全局归一化因子做近似时，可以用随机特征或频域核展开的方法估计核函数，这本质上和 FFT 的频域视角同源。所以我在任何场合都会强调：FFT 优化 softmax 的姿势不是直接套 FFT，而是先优化喂给 softmax 的那一段计算链路。
4.5 融合优化的整体流程
把 attention 相关的算子放在一起看，一个合理的 FFT 融合优化流程是这样：
import numpy as np def fft_attention(Q, K, V, pos_bias, d): # Q, K: [N, d] 张量，pos_bias: [N, N] 或 Toeplitz 形式 # 1. 对相对位置偏置部分用 FFT 加速 bias_term = np.fft.ifft( np.fft.fft(pos_bias_row, axis=0) * np.fft.fft(np.ones((N, 1)), axis=0) ).real # 示意代码，简化表达 # 2. 内容相关部分保留矩阵乘 score_content = Q @ K.T / np.sqrt(d) # 3. 拼接并 softmax scores = score_content + bias_term probs = softmax(scores, axis=-1) return probs @ V注意这段代码只是思路示意，真实实现要处理 batch、多 head、复数 layout，以及 FFT 结果的共轭对称问题。工程上的关键点在于：把能走 FFT 的部分切得足够干净，不要让非结构计算拖住整体性能；同时注意 softmax 前的数据格式转换成本，一次 layout 转换可能吃掉 FFT 省下的所有时间。
5. 一个混合优化示例的完整复盘
5.1 测试平台与基准方法
为了验证"FFT 优化其他 kernel 和算子"这件事到底值不值，我在一个自研推理引擎的小型基准集上做了一轮完整对比。平台是 16 核 AVX512 的服务器 CPU，数据精度 float32，基准实现全部用本地调优过的常规方案（im2col+GEMM、手写 softmax、标准 attention 实现），FFT 使用经过对齐优化的自研 radix-2 实现，没有依赖外部库。
测试集的算子覆盖：3×3 卷积、5×5 卷积、7×7 卷积、稠密 GEMM 512×512、Toeplitz 结构化 GEMM 1024×1024、序列长度 4096 的 attention、纯 softmax 4096×4096。每个算子跑 100 次取中位数。
5.2 每个算子的优化前后对比
| 算子 | 基线耗时 | FFT 优化后 | 加速比 | 备注 |
|---|---|---|---|---|
| 3×3 卷积 112×112 | 3.10 ms | 4.12 ms | 0.75x | FFT 反而更慢 |
| 5×5 卷积 112×112 | 7.82 ms | 5.61 ms | 1.39x | padding 后 FFT 收益为正 |
| 7×7 卷积 112×112 | 14.55 ms | 8.94 ms | 1.63x | 大核优势明显 |
| 稠密 GEMM 512×512 | 0.42 ms | 0.68 ms | 0.62x | 无结构，FFT 亏 |
| Toeplitz GEMM 1024×1024 | 2.31 ms | 0.77 ms | 3.00x | 结构矩阵吃满 FFT 红利 |
| Attention 4096 | 61.2 ms | 44.8 ms | 1.37x | 主要收益来自位置 bias 部分 |
| Softmax 4096×4096 | 2.10 ms | 2.30 ms | 0.91x | FFT 引入额外搬移，略亏 |
5.3 优化结果的复盘解读
这组数据很典型，也很诚实。FFT 不是在所有地方都赢。3×3 卷积亏，稠密 GEMM 亏，纯 softmax 也亏；但大卷积核赚、Toeplitz 结构矩阵赚、带相对位置编码的 attention 赚。把这些放在一起看，结论很清晰：FFT 优化的核心是"找结构"，结构越明显，FFT 越能发挥O(N log N)对O(N²)的复杂度碾压。
还有一个观察：attention 的加速比 1.37x 看起来不够惊艳，但它的意义在于没有改动模型结构的前提下，从 attention 内部的数学结构里挤出性能。如果配合频域 token mixing 的替代方案，加速比会完全不同，但那就属于模型结构的改造了，不是纯算子优化。两者要分清。
6. 常见问题与排查技巧实录
6.1 问题速查表
| 现象 | 可能原因 | 解决方法 |
|---|---|---|
| 输出边缘像素明显错误 | 循环卷积混叠，未 padding 到线性卷积长度 | 按 N+M-1 补齐，或使用重叠保留法 |
| FFT 后性能反而更慢 | 算子本来 memory-bound，或核尺寸太小 | 换 Winograd / 常规 GEMM，别硬套 FFT |
| 结果精度波动大 | 输入动态范围大，FFT 相位信息丢失 | 适当缩放输入，或对敏感层保持时域计算 |
| 复数内存占用翻倍后频繁 cache miss | 复数布局未对齐，FFT 输出立即被拷贝 | 使用 in-place FFT 并提前分配连续内存 |
| 多通道卷积里卷积核重复变换 | 每个 batch 都在做 kernel FFT | 把 kernel 的 FFT 结果缓存复用 |
| 频域点乘后 shape 对不上 | 实部/虚部 layout 处理错误 | 明确库的复数排布是 RCC 还是 CCS |
6.2 精度问题的定位方法
FFT 引入的数值误差通常是可控的，但一旦出现精度下降，定位思路要清晰。第一步，关掉所有 FFT 路径，确认基线精度；第二步，只对某一层开启 FFT，逐层二分定位是哪一层出的问题；第三步，检查该层的输入动态范围，如果输入数值跨度过大（比如大于 1e4），FFT 的加减运算会先丢失小数值的尾数位，这时可以在时域做一次全局缩放，或者改用 double 做 FFT 后再转回 float。
我碰到过的另一种精度坑是频域点乘后直接 IFFT，中间把实部和虚部当成独立通道做了 ReLU 之类的非线性操作，这相当于在频域强行截断相位信息，恢复出来的时域信号会产生振铃效应，输出看起来就是"脏脏的条纹"。记住，频域数据只有在做线性滤波时才是安全的，任何非线性操作都应回到时域再做。
6.3 性能和内存的实用建议
FFT 优化算子的工程实现里，内存规划经常决定成败。首先，尽量使用 in-place FFT，减少一次完整复数矩阵的拷贝；其次，把复数数据按通道连续排布，用 n 个独立 FFT 而不是 n 次单点循环调用；第三，在 CPU 上做批量小 FFT 时，优先利用库的 batch 接口而不是自己写循环，后者会因为函数调用开销损失两到三成性能。
还有一条针对多算子融合的经验：FFT 的输出不要急着 IFFT 回时域，如果下一步仍然是线性算子（比如另一个卷积或者矩阵乘），完全可以在频域连续做多次操作，最后只做一次 IFFT。这在 attention 链路里尤其好用，QK^T 和 V 聚合都是线性运算，理论上可以全部压在频域完成，省掉两次时频切换。
6.4 什么时候该果断放弃 FFT
我做优化经常提醒自己：FFT 不是信仰，它是一种工具。如果你遇到下面任意一种情况，就果断放弃，别恋战：
- 单次算子规模小（矩阵维度小于 256、卷积核小于 3×3、序列长度小于 512）。
- 数据是纯逐元素操作，比如激活函数、归一化、softmax 本身。
- 硬件上的 FFT 库性能极差且没有替代实现。
- 数值调试已经超过 3 天还没有收敛。
在上述条件下，把时间花在手写 SIMD、算子融合、内存布局优化上，回报通常更大。我的一个项目里本来计划用 FFT 优化所有卷积，实测完发现只有大核卷积和 attention 值得改造，于是砍掉大部分 FFT 路径，把节省的时间投入到通用 GEMM 优化上，最终整体推理速度反而提升更多。这个教训让我确立了一个习惯：任何优化手段都要先用小实验验证盈亏平衡点，再决定是否规模化推广。
我在实际调优里还有一个体会：FFT 优化的收益常常不是单算子层面的，而是整个计算图层面的。当你把某个底层算子换成频域实现后，周围的 layout 转换、kernel 融合策略、内存复用方案都会跟着变。我见过一个反例，某同事花大力气把单个卷积换成了 FFT，结果因为前后算子都要把数据转 f32，整个网络反而变慢。所以接入 FFT 之前，先画一下数据流转路径，把每个转换点都算进开销里，再做决策。
最后分享一个小技巧，也是我后面每次都会先做的事：在做 FFT 优化之前，先把目标算子用 profiler 拉一遍，看它的算术强度和访存占比。如果算术强度已经超过某个阈值（比如在 CPU 上大于 32 flops/byte），再考虑 FFT；否则，优化方向应该先指向减少数据搬移。这个判断十次有九次能帮你少走弯路。
