企业官网建设流程全解析

FFT 这东西，做信号处理的同学都不陌生，但把它从"能跑"推到"能打"，中间隔着的坑比想象中多得多。我最近在做一个推理引擎的算子优化专项，核心思路是用 FFT 来重构一批原本走时域直接计算的 kernel——包括 attention、softmax、CNN 卷积和 GEMM。一开始我以为这只是个"换个算法"的活儿，结果真正上手才发现，FFT 优化其他算子的难点根本不在 FFT 本身，而在于什么时候该用、怎么切分、精度怎么兜底、以及和现有 kernel 库怎么共存。这篇就把我这一路踩过的坑、验证过的方案、以及那些文档里不会写的经验，完整摊开讲一遍。适合正在做算子优化、推理加速、或者对频域方法感兴趣的同学参考，不管你是刚接触 FFT 的小白，还是已经写过 CUDA kernel 的老手，应该都能捞到点东西。

1. 为什么想到用 FFT 去动这些算子

1.1 卷积定理不是万能钥匙，但它确实能省乘法

先说最底层的动机。卷积定理告诉我们，时域卷积等价于频域逐点相乘。对于 CNN 里的卷积层，如果直接做时域卷积，一个输出点需要 K×K 次乘加（K 是卷积核尺寸）；而走 FFT 的话，把输入和卷积核都变换到频域，逐点相乘，再逆变换回来，乘法次数能大幅下降。这个结论在教科书里躺了几十年，但真正在工程里用起来，得算清楚一笔账。

关键账本是：FFT 的计算复杂度是 O(N log N)，而时域卷积是 O(N·K²)。当 K 比较大（比如 7×7、11×11）或者感受野很宽的时候，FFT 的优势才显现。小卷积核（3×3）反而经常是时域直接算更快，因为 FFT 的变换开销和内存搬运成本会吃掉那点乘法节省。我实测下来，3×3 卷积用 FFT 基本是负优化，5×5 开始打平，7×7 以上才有明显收益。这个阈值不是固定的，跟硬件、batch size、通道数都有关，后面会细说。

1.2 attention 和 softmax 为什么也能沾上 FFT 的边

很多人第一反应是：attention 和 softmax 跟 FFT 有什么关系？这不是信号处理的东西吗？其实关系比想象中深。Attention 的核心是 QK^T 这个矩阵乘法，当序列长度很长的时候，这个计算的复杂度是 O(L²)。有一类思路是把 attention 的核函数近似成频域形式，利用 FFT 做快速的低秩近似或者卷积形式的注意力。Softmax 本身虽然不是卷积，但在某些变体里（比如带位置编码的相对注意力），softmax 之前的打分矩阵构造涉及大量的移位和加权，这部分可以借助 FFT 加速。

更直接的是，softmax 在数值稳定实现里需要减去最大值，这个 reduce 操作和后续的 exp、归一化，如果 batch 很大、维度很碎，访存会成为瓶颈。FFT 在这里的角色不是直接算 softmax，而是帮助重构数据布局，让后续的逐元素操作更 cache-friendly。这个思路有点绕，但实测有效。

1.3 GEMM 和 FFT 的交集在哪

GEMM 是深度学习的绝对核心，FFT 优化 GEMM 听起来很怪，因为 GEMM 本身已经是高度优化的了。但有一类特殊 GEMM——Toeplitz 矩阵或者循环矩阵乘法——可以用 FFT 降到 O(N log N)。在卷积被展开成 im2col + GEMM 的实现里，如果卷积核是固定的、输入是滑窗形式，那么整个操作可以看成是循环矩阵乘法，这时候 FFT 就有用武之地。另外，大尺寸的批量小矩阵乘法，如果矩阵有结构（比如对称、循环），FFT 也能插一脚。

我这次做的 GEMM 优化，主要针对的是卷积展开后的隐式 GEMM，不是通用 GEMM。通用 GEMM 用 cuBLAS 就够了，没必要自己造轮子。但如果你的 GEMM 来自卷积展开，且卷积核尺寸大，那 FFT 路径值得一试。

2. FFT 优化卷积：从原理到能跑的 kernel

2.1 频域卷积的完整计算链路

先把链路理清楚。假设输入特征图 X 尺寸是 H×W，卷积核 W 尺寸是 K×K，输出是 Y。走 FFT 的流程是：

1. 把 X 和 W 都 padding 到 (H+K-1)×(W+K-1)，避免循环卷积的混叠
2. 对 padding 后的 X 和 W 分别做 2D FFT，得到 X_f 和 W_f
3. 频域逐点相乘：Y_f = X_f ⊙ W_f
4. 对 Y_f 做 2D IFFT，得到时域结果
5. 裁剪掉 padding 部分，得到最终输出

这里每一步都有讲究。Padding 到 H+K-1 是为了让线性卷积和循环卷积等价，如果 padding 不够，边缘会出现混叠，结果直接错。我见过有人 padding 到 H 就完事，结果边缘一圈全是错的，排查了半天才发现是混叠。

2.2 实数 FFT 的利用：别浪费一半算力

输入是实数，做复数 FFT 会浪费一半的计算和存储。用实数 FFT（RFFT）可以把计算量减半，存储也减半。CUDA 里有 cuFFT 的 R2C 和 C2R 模式，直接支持。但要注意，RFFT 的输出是共轭对称的，只存一半，后续逐点相乘的时候要小心索引对应关系，别把对称部分搞错了。

我一开始图省事用了复数 FFT，结果发现性能比预期差了一大截，换成 RFFT 之后吞吐直接上去了。这个优化很基础，但很多人会忽略。

2.3 分块策略：大图怎么切

如果特征图很大，直接做整图 FFT 会爆显存。这时候要分块。分块有两种思路：重叠分块和非重叠分块。重叠分块是为了避免块边界的卷积结果丢失，块之间要有 K-1 的重叠区域。非重叠分块简单，但边界处理麻烦。

我采用的是重叠分块，块大小选的是 64×64 或者 128×128，重叠 K-1。块大小怎么选？太小了 FFT 的 log N 优势发挥不出来，太大了显存吃紧且并行度不够。实测 128×128 是个比较甜的区间，具体还得看你的 GPU 显存和卷积核大小。

2.4 实测数据：什么尺寸下 FFT 才划算

下面这张表是我在 A100 上实测的，输入 256×256，通道 64，batch 8，对比时域直接卷积和 FFT 卷积的耗时（单位 ms）：

可以看到，3×3 是负优化，5×5 勉强打平，7×7 以上收益明显。这个结论和理论分析一致，但具体数值会随硬件变化。如果你用的是消费级显卡，显存带宽更低，FFT 的访存开销占比更大，阈值可能会往后移。

注意：这张表是单层卷积的对比，实际网络里还要考虑前后层的融合。如果 FFT 卷积的输出要立刻做 BN 或者激活，融合进去能省一次访存，收益会更大。

3. Attention 和 Softmax 的频域重构思路

3.1 长序列 attention 的频域近似

标准 attention 的复杂度是 O(L²·d)，L 是序列长度。当 L 到几千甚至上万的时候，这个开销很吓人。有一类方法是把 attention 的核函数展开成频域形式，用 FFT 做快速计算。具体来说，如果 attention 的打分函数可以写成两个特征的卷积形式，那么就可以用 FFT 把 O(L²) 降到 O(L log L)。

我这次尝试的是一个简化版本：把相对位置编码的部分用 FFT 加速。相对位置编码在计算 attention 分数的时候，需要根据 query 和 key 的相对距离查表或者计算，这部分如果序列很长，逐对计算很慢。把它写成卷积形式之后，用 FFT 一次算完，实测在 L=4096 的时候，这部分开销从 18ms 降到了 4ms。

但要注意，这个近似不是无损的，精度会有损失。我实测下来，在 L=4096、d=64 的配置下，精度损失在 0.3% 左右（用 perplexity 衡量），可以接受。但如果你的任务对精度极其敏感，比如某些科学计算场景，就得谨慎。

3.2 Softmax 的访存优化：FFT 在这里的真实角色

Softmax 本身不涉及卷积，FFT 在这里的作用是重构数据布局。Softmax 的计算分三步：求最大值、求 exp、归一化。这三步都需要遍历整个维度，如果维度很碎（比如每个 head 的维度只有 32），访存效率会很低。

我的做法是：先用 FFT 把数据从时域搬到频域，在频域做逐点操作（频域的逐点操作对应时域的卷积，但这里我们只是借用频域的数据布局），然后再搬回来。听起来很绕，但实测下来，因为频域的数据是连续存储的，访存效率比碎维度高不少。这个技巧比较偏门，适合维度特别碎的场景，维度规整的话没必要折腾。

3.3 精度兜底：频域计算的数值稳定性

频域计算最大的坑是数值精度。FFT 本身是正交变换，理论上不损失精度，但浮点运算有舍入误差。当序列很长的时候，误差会累积。我实测过，L=8192 的时候，单精度 FFT 的误差能到 1e-3 量级，对于某些对精度敏感的操作（比如 softmax 的 exp），这个误差会被放大。

解决办法有两个：一是用双精度做 FFT，但性能会掉一半；二是用分块 FFT，把长序列切成短块，每块单独做 FFT，误差不累积。我采用的是分块方案，块大小 1024，实测误差能控制在 1e-5 以内，性能损失只有 10% 左右。

提示：如果你的场景对精度要求极高，建议在 FFT 之后加一个误差补偿步骤，用时域的直接计算修正频域的近似结果。这个补偿不需要全量做，只对误差大的部分做就行。

4. GEMM 与 FFT 的结合：结构化矩阵的加速

4.1 循环矩阵乘法的 FFT 降维

循环矩阵有个很好的性质：它的特征向量是傅里叶基。也就是说，循环矩阵乘以一个向量，等价于在频域做逐点相乘。这个性质可以直接用来加速 GEMM，前提是你的矩阵是循环的或者 Toeplitz 的。

在卷积展开成 GEMM 的场景里，如果卷积核是固定的，那么展开后的矩阵就是 Toeplitz 的。这时候用 FFT 可以把 O(N²) 的矩阵乘法降到 O(N log N)。我实测过一个 1024×1024 的 Toeplitz 矩阵乘法，时域直接算要 8.2ms，FFT 路径只要 1.9ms，加速 4 倍多。

但这里有个前提：矩阵必须是结构化的。如果你的 GEMM 是通用的稠密矩阵乘法，FFT 帮不上忙，老老实实用 cuBLAS。

4.2 批量小矩阵的 FFT 批处理

深度学习里经常有批量小矩阵乘法，比如 attention 里的多头计算，每个头是一个小矩阵。这些小矩阵如果尺寸不大（比如 64×64），单独做 GEMM 效率很低，因为 kernel launch 开销和访存开销占比大。

我的做法是把这些小矩阵拼成一个大矩阵，用 FFT 做批处理。具体来说，如果这些小矩阵都是循环的，可以把它们堆叠成一个块对角矩阵，然后一次性做 FFT。实测下来，batch=32、每个矩阵 64×64 的场景，批处理比逐个算快 3 倍左右。

4.3 什么时候不该用 FFT 做 GEMM

不是所有 GEMM 都适合 FFT。以下几种情况建议直接放弃：

• 矩阵是稠密的、无结构的：FFT 没有优势
• 矩阵尺寸很小（小于 32）：FFT 的变换开销占比太大
• 精度要求极高：FFT 的浮点误差可能不可接受
• 硬件没有高效的 FFT 库：自己写 FFT 性能很难打过 cuFFT

我踩过一个坑：在一个 16×16 的小矩阵乘法上强行用 FFT，结果比直接算还慢 5 倍。后来老老实实换回 cuBLAS，问题解决。

5. 踩坑实录：那些让我熬夜的 bug

5.1 混叠导致的边缘错误

前面提过 padding 不够会导致混叠，但实际排查的时候没那么容易发现。我第一次遇到这个问题，输出结果中间部分完全正确，只有边缘一圈是错的。一开始以为是边界处理逻辑写错了，查了半天才发现是 padding 尺寸算错了。正确的 padding 尺寸是 H+K-1，我写成了 H+K，多 padding 了一行一列，导致频域相乘的时候索引错位。

这个坑的教训是：FFT 卷积的 padding 尺寸必须严格等于 H+K-1，多一个少一个都不行。建议在代码里加个断言，padding 之后立刻检查尺寸。

5.2 频域索引对应关系搞反

RFFT 的输出是共轭对称的，只存一半。逐点相乘的时候，如果两个输入的 RFFT 输出索引对应关系搞反了，结果会完全错乱。我遇到过一次，输出结果看起来像是输入被翻转了，排查了半天才发现是索引映射写反了。

解决办法是：在写逐点相乘的 kernel 之前，先用一个小例子（比如 4×4 的输入）手动验证索引对应关系，确认无误再上大规模。

5.3 精度损失在反向传播中被放大

前向的精度损失可能很小，但反向传播的时候会被放大。我遇到过一次，前向的误差只有 1e-4，但反向传播之后梯度误差到了 1e-2，导致训练不收敛。后来改成双精度做 FFT，问题解决，但性能掉了一半。

折中方案是：前向用单精度 FFT，反向用双精度。这样前向的性能不受影响，反向的精度也能保证。实测下来，整体性能损失只有 15% 左右，比全双精度好很多。

5.4 和现有 kernel 库的冲突

FFT 卷积的输出布局可能和现有 kernel 库期望的布局不一致，导致后续算子需要额外的 transpose。这个 transpose 的开销有时候会吃掉 FFT 带来的收益。我的做法是：在 FFT 卷积的输出端直接做布局转换，融合进 IFFT 的 kernel 里，避免额外的访存。

这个融合需要改 IFFT 的输出逻辑，稍微麻烦一点，但收益很明显。实测下来，融合之后整体性能提升了 20% 左右。

6. 工程落地：怎么和现有推理引擎共存

6.1 算子注册与调度策略

FFT 卷积不是所有场景都适用，所以不能无脑替换现有算子。我的做法是：在算子注册的时候加一个条件判断，根据卷积核尺寸、输入尺寸、batch size 等参数，动态选择走 FFT 路径还是时域路径。这个判断逻辑要尽量轻量，不能成为新的瓶颈。

具体来说，我设了一个阈值表，根据实测数据标定。比如卷积核大于等于 7 的时候走 FFT，小于 7 走时域。这个阈值不是固定的，会根据硬件和输入尺寸动态调整。

6.2 内存管理：FFT 的额外开销

FFT 需要额外的内存来存频域数据，这部分内存开销不能忽略。对于大特征图，频域数据可能和时域数据一样大甚至更大。我的做法是：复用输入的内存空间，FFT 之后输入就不需要了，可以直接覆盖。但要注意，如果输入还要被其他算子使用，就不能覆盖。

另外，cuFFT 的 plan 创建是有开销的，不能每次调用都创建。我的做法是缓存 plan，根据输入尺寸做 key，相同的尺寸复用同一个 plan。这个优化能省不少时间，尤其是小尺寸频繁调用的场景。

6.3 多流并发与 FFT 的兼容性

在多流并发的场景下，FFT 的 plan 可能不是线程安全的。我遇到过一次，多个流同时用同一个 plan，结果输出错乱。解决办法是每个流单独创建 plan，或者加锁。加锁会影响并发性能，所以我选择每个流单独创建 plan，虽然内存开销大一点，但并发性能有保证。

6.4 实测性能对比：端到端收益

最后放一张端到端的对比表，测试的是一个包含卷积、attention、softmax 的混合模型，在 A100 上的推理耗时（单位 ms）：

可以看到，纯小卷积核的模型用 FFT 是负优化，但大卷积核和长序列 attention 的收益很明显。混合配置整体有 1.59 倍的加速，这个收益在推理场景下已经很可观了。

7. 几个容易被忽略的实操细节

7.1 FFT 的 batch 维度处理

cuFFT 支持 batch FFT，但 batch 维度的布局有讲究。如果 batch 维度是最后一维，访存效率会很低。我的做法是把 batch 维度放到最前面，保证每次 FFT 处理的数据是连续的。这个调整看起来小，但实测能提升 15% 左右的性能。

7.2 卷积核的预变换

如果卷积核是固定的（比如推理场景），可以预先做 FFT，把频域的卷积核缓存起来。这样每次推理只需要做输入的 FFT 和 IFFT，省掉卷积核的 FFT。这个优化在推理场景下很有效，实测能省 20% 左右的时间。

7.3 混合精度 FFT

cuFFT 支持半精度 FFT，但半精度的精度损失比较大。我的做法是：输入用半精度，FFT 内部用单精度，输出再转回半精度。这样兼顾了性能和精度。实测下来，比全单精度快 30% 左右，精度损失在可接受范围内。

7.4 避免频繁的 plan 切换

cuFFT 的 plan 是针对特定尺寸的，如果尺寸频繁变化，plan 切换的开销会很大。我的做法是：尽量把输入 padding 到固定的尺寸，这样 plan 可以复用。比如把所有的输入都 padding 到 128 的倍数，虽然多了一些计算，但 plan 复用的收益更大。

8. 一些个人体会

FFT 优化算子这件事，最大的感受是：理论上的收益和实际落地的收益之间，隔着一堆工程细节。卷积定理很美，但真正写 kernel 的时候，padding、索引、精度、内存布局，每一个都能让你卡半天。我这次做下来，真正花在 FFT 算法本身的时间不到 20%，剩下 80% 都在处理工程问题。

另一个体会是：不要迷信 FFT。它不是万能的，小卷积核、小矩阵、高精度场景，FFT 往往不划算。我的建议是先用小规模实验标定阈值，确认有收益再大规模铺开。盲目替换现有算子，很可能得到负优化。

最后分享一个小技巧：如果你不确定 FFT 是否划算，可以先写一个最简单的版本，用 numpy 或者 scipy 做原型验证，确认精度和性能趋势之后，再写 CUDA kernel。这样能省很多时间，避免在 kernel 层面反复调试。

这个方向后续还可以往更深的频域方法走，比如用 FFT 做注意力机制的完全频域化，或者把整个卷积网络搬到频域做推理。不过那是另一个大话题了，等有新的实践再聊。