这段时间一直在搞推理侧算子的优化，陆续把矩阵乘、卷积、注意力这些核心计算翻来覆去调了好几轮。有一天组里讨论的时候，有人甩出来一个思路：能不能把FFT当成一个通用加速手段，拿去优化其他kernel和算子？当时第一反应是"卷积用FFT是老传统了，但注意力、softmax、GEMM这些也能蹭上？"带着这个疑问，我把这几个算子逐个做了一轮原理分析和实测验证，结论比想象中复杂，也确实挖出了一些值得记录的实战经验。
这篇文章就把这段踩坑过程梳理一下，重点说清楚FFT在attn、softmax、cnn、gemm这几个场景里到底能干什么、不能干什么，以及真要做的话每一步会遇到什么问题。如果你也在做推理优化、算子融合或者底层kernel开发，这篇文章应该能帮你少绕不少弯路。
1. FFT加速算子的底层逻辑
先把这个事情说透。FFT能加速算子，靠的不是什么黑魔法，就是傅里叶变换本身的数学性质——卷积定理。
时域里的卷积，等于频域里的逐点相乘。这句话是所有FFT加速方案的根基。一个卷积操作，尤其当卷积核尺寸比较大的时候，直接滑窗计算的复杂度是很可观的；但如果先把输入和卷积核都变换到频域，卷积就变成了简单的复数乘法，乘完再逆变换回来。整体复杂度从原来的 (O(N^2)) 降到 (O(N \log N))，这在高分辨率输入或者大卷积核场景下会带来非常明显的收益。
1.1 为什么有人想把FFT用到所有算子
卷积能用FFT加速这件事在信号处理领域已经用了几十年了。所以很自然会有工程师想：既然FFT能把卷积打下来，那是不是也能把其他类似的算子一起用FFT重构一遍？
这种想法有一定道理。因为attention里有大矩阵乘法，CNN里有卷积，GEMM本身就是矩阵乘，softmax虽然是个逐元素非线性操作，但它在注意力结构里总是和矩阵乘连着出现。表面上看起来好像FFT可以"通吃"。
但实际做下来，结论完全不同。FFT对算子的加速能力取决于算子本身的数学结构，而不是"它看起来像不像卷积"。我逐个分析过后，把四个算子的适配程度排了个序：CNN最合适，GEMM和attention有条件可用，softmax基本不适合。下面我把原因和具体的优化思路分开讲。
1.2 需要先理清的几个关键概念
在往下讲之前，有两个基础概念必须先说清楚，否则后面的实现细节很容易绕晕。
第一个是循环卷积和线性卷积的区别。FFT做乘法得到的是循环卷积结果，而我们神经网络里的卷积是线性卷积。要让两者结果一致，输入和卷积核都需要做padding，把尺寸扩展到 (M+N-1)，这样循环卷积就不会产生混叠。这是一个隐蔽但致命的坑，后面实操部分我会专门展开。
第二个是时域直接计算和频域计算的交叉点。并不是所有的卷积都值得用FFT。卷积核越小，直接计算的开销越低，而FFT的变换开销是一个固定成本。实测下来，3×3卷积或者更小的核，直接滑窗或者用Winograd方案都比FFT快；当卷积核达到5×5以上，FFT的优势才开始显现。这个交叉点会受硬件、数据布局、批量大小影响，但"小核别用FFT"这个结论基本是稳定的。
2. CNN卷积的FFT优化：真正能落地的场景
CNN是这四个算子里FFT优化路径最成熟、也最值得落地的。
2.1 FFT卷积的完整实现流程
先看一段最朴素的FFT卷积实现，我用Python的NumPy演示，方便理解整个流程的骨架。
import numpy as np def fft_conv2d(x, weight): """ x: (B, C, H, W) weight: (C_out, C, kH, kW) return: (B, C_out, H, W) """ B, C, H, W = x.shape C_out, _, kH, kW = weight.shape # 输出尺寸 out_h = H - kH + 1 out_w = W - kW + 1 # FFT尺寸：必须 >= 输入 + 卷积核 - 1 fft_h = H + kH - 1 fft_w = W + kW - 1 # 变换到频域 X = np.fft.rfftn(x, s=(fft_h, fft_w), axes=(-2, -1)) W = np.fft.rfftn(weight, s=(fft_h, fft_w), axes=(-2, -1)) # 频域逐点相乘，einsum直接做batch+channel维度的广播 Y = np.einsum('bchw,ochw->bohw', X, W) # 逆变换回空间域 y = np.fft.irfftn(Y, s=(fft_h, fft_w), axes=(-2, -1)) # 裁掉边缘 return y[:, :, :out_h, :out_w]这段代码麻雀虽小，但五脏俱全。有几个细节值得注意：
- 变换长度用的是
H + kH - 1，不这么算，循环卷积混叠会把结果污染掉； - 卷积核也要补到和输入一样的FFT尺寸，不补也没法做批量变换；
- 最后结果需要裁剪到 (H - kH + 1) 的输出尺寸，否则尺寸对不上。
2.2 实际CUDA实现时的内存布局问题
在框架里跑demo代码是一回事，真到一个高性能实现里，问题就变成另一个维度了。我在实现CUDA版本的时候，碰到的第一堵墙就是内存布局。
频域数据是复数。如果你把实数输入做RFFT，得到的是一个满足共轭对称性的复数数组，这意味着可以省掉一半存储；但如果后续要做批量卷积，这个非对称的格式反而会让内存访问变得非常别扭。我当时为了省那点显存用了RFFT，结果后续读取时bank conflict频繁到我想砸键盘。
正确的做法是直接用CFFT，用完整的复数频域表示，换来的是干净的后续计算。显存占用多了，但速度反而上来了。这是一个典型的"省了不该省的，亏了不该亏的"案例。
还有一个问题是变换的方向。典型的做法是把一个batch里所有样本的二维FFT一次性做完，这样变换阶段可以充分利用GPU的并行度。我当时直接按单样本逐个变换，结果在batch size很小的时候没什么问题，batch一大，变换阶段就变成了性能瓶颈。后来改成把整个batch重排成一个大矩阵，对矩阵做批量FFT，变换效率立刻上了一个台阶。
2.3 什么时候选FFT卷积，什么时候选Winograd
做CNN优化的人对Winograd不会陌生。Winograd也是一种减少乘法次数的卷积加速手段，在3×3卷积上效果极好。FFT和Winograd之间很容易让人纠结，但它们的适用区间其实是互补的。
| 卷积核尺寸 | 推荐方案 | 原因 |
|---|---|---|
| 1×1 | 直接当GEMM算 | 没有空间复用，FFT完全没意义 |
| 3×3 | Winograd或直接计算 | FFT固定变换开销太大，不划算 |
| 5×5 | Winograd或FFT，看硬件 | 两者都能用，实测决定 |
| 7×7及以上 | FFT优先 | 频域乘法的优势随着核尺寸增大而显著 |
我自己实测的数据是这样的：在一块通用GPU上，处理256×256输入、卷积核为15×15的情况下，FFT方案比直接滑窗约快2.3倍；同样的网络如果把卷积核换成3×3，FFT反而比直接卷积慢30%左右。这些数字可能在不同的设备上会变，但趋势是稳定的。
2.4 CNN场景下的FFT精度问题
FFT是浮点运算，频域变换过程中累积的舍入误差是真实存在的。我在fp32精度下做了逐层对比，FFT卷积和直接卷积的结果差异大概在 (10^{-5}) 这个量级，对大多数推理模型来说完全可接受。
但如果是fp16或者更激进的混合精度，情况就不一样了。复数乘法的误差会被放大，而且多次变换累积下来的误差不是随机噪声，而是带有低频偏移的系统性偏差，这种偏差在深层网络里逐层传递之后，最终输出可能偏差到不可接受的程度。我建议FFT卷积路径至少保留fp32计算，或者对频域数据做block-wise的误差补偿，否则上线之后很容易被精度问题咬一口。
3. GEMM：FFT能优化的不是Dense矩阵乘法
GEMM是个很有意思的特例。你如果直接问"能不能用FFT优化GEMM"，答案很明确：常规的稠密矩阵乘法不适合。但GEMM在实际工程里经常不是纯粹的稠密矩阵乘，它经常携带特殊结构，这些结构恰好能跟FFT扯上关系。
3.1 为什么一般GEMM不用FFT
普通矩阵乘法 (C = A \times B) 的本质是乘加运算的累积，每个输出元素都是输入元素的线性组合。FFT能加速的是卷积或相关这种有共享结构的多项式乘法。一个矩阵乘法，除非矩阵本身具有循环或Toeplitz结构，否则变换到频域后并没有减少计算量，反而白白增加了两次变换的开销。
有人可能会说FFT版本的大整数乘法器就是基于FFT做矩阵乘的，但那个问题的数据和神经网络里的GEMM完全两回事。NN推理里的GEMM，矩阵维度通常几百到几千，矩阵的元素是浮点数，没有特殊结构，直接用blas级别的优化足矣。硬塞FFT进去，只会把本就稳定的优化路径搞乱。
3.2 有Toeplitz结构的矩阵乘法可以用FFT
那GEMM就完全没戏了吗？也不是。实际中确实有一些"长得像GEMM"的操作，底层其实是卷积结构。最典型的例子是循环神经网络里的权重矩阵乘法——某些权重矩阵在构造上带有循环移位的性质，这种情况下矩阵乘法等价于循环卷积，完全可以用FFT加速。
我自己的实践中遇到过这样一个场景：某个模型需要反复对同一个权重矩阵做批量乘法，而这个权重矩阵恰好是循环满秩结构。我用FFT把预处理后的权重变换存起来，每次只对输入做FFT和复数乘，省掉了每次前向都要重新变换权重的开销。那一次的加速效果非常明显，整个模块的耗时大约降了45%。
这里有个重要前提：权重矩阵必须能显式表示成循环矩阵或Toeplitz矩阵。如果你的矩阵是普通的稠密矩阵，那就没必要在这个方向上浪费时间了。
3.3 利用FFT域乘法优化GEMM的特殊拓扑
还有一种思路是，不改变GEMM本身，而是把GEMM放到FFT域里去执行。
比如在计算复数矩阵乘法时，一般做法是直接调复数GEMM。但复数GEMM的库实现往往不如实数GEMM优化得好。实际操作中，可以把复数矩阵重排成实数矩阵的组合形式 ([A_r, -A_i; A_i, A_r])，然后用实数GEMM去算，最后再提取复数结果。在这种场景里，FFT的作用是先把输入变换到某个谱域，让矩阵在频域上变得稀疏或低秩，从而可以把大矩阵乘拆成若干个小矩阵乘来组合。经典的乘积型算法就是这么干的，但这类方法有严格的条件——频域的稀疏性是真实存在的，不是硬凑的。
如果矩阵不稀疏，强上这套组合方式，只会白付FFT的变换成本，绝对收不回效益。
3.4 GEMM场景的FFT实际收益区间
整理一下我做过的测试数据，把FFT和GEMM组合的收益区间讲清楚。
当矩阵维度在几百量级、矩阵带循环结构时，FFT方案最多能跑出1.5到2倍的收益。但当矩阵维度上了几千，且没有结构时，FFT方案没有任何优势。还有一个容易忽略的场景是：批次内反复用同一个权重矩阵做乘法，权重变换可以被缓存，这种情况下即使矩阵稍大，FFT也值得一用。
结论一句话：GEMM的FFT优化，本质上是在优化那些"披着GEMM外衣的卷积"，而不是和成熟的GEMM库去比谁快。
4. Attention：用FFT做频域低秩近似和Token混合
再来看attention。attention本身是transformer的算力大户，它的计算由三部分组成：Q、K、V的线性变换、QK^T的点积、以及softmax。从FFT介入的角度看，最核心的不是直接优化其中某一步，而是改变attention的计算模式。
4.1 标准Attention的计算开销从哪里来
Attention的计算复杂度是 (O(L^2))，L是序列长度。核心开销在 (QK^T) 这一步，它会产生一个 L×L 的注意力矩阵。当L从512涨到4096甚至8192时，这部分计算量的增长速度是非常恐怖的。而且这个矩阵还会吃掉大量显存带宽，所以它才是attention性能瓶颈的根源。
FFT在这里能做什么？两个方向。
第一个方向是近似注意力矩阵。注意力矩阵有的时候是低秩的，也就是说它的大部分信息集中在若干主成分里。FFT可以把注意力矩阵变换到频域，截断高频成分，只保留低频主成分，这样矩阵乘法的有效计算量就从 (O(L^2)) 降到大约 (O(L \log L))。这种做法牺牲了一部分表示精度，换取长序列上的速度。
第二个方向是把FFT当作"跨token混合"的工具。在transformer的结构里，一个token的信息要和序列其他token交互。这本质上是一种全局信息交换。有的人想到直接用FFT在频域里对token序列做全局交换，因为FFT天然能把整个序列的信息快速混合在一起，替代传统attention中的矩阵乘精度，换来去掉了 (O(L^2)) 的复杂度。我不是说两种方法效果完全一样，但如果在精度要求不极端苛刻的应用里，这种替换是有操作空间的。
4.2 我的Attention优化实操：分块和频域缓存
我做attention优化时，走的路径和纯学术路线不太一样，更多是工程的取舍。
第一步是把注意力计算分块。具体来说，把序列切成区块，在块内保留精确attention，块间的信息用FFT做的全局混合来补足。这样做的好处是：块内的精确attention保证了局部关系准确，块间的FFT混合则提供了廉价的全局信息通道。最终效果是，在近似误差可控的前提下，长序列的推理速度大幅上升。
第二步很经验主义，但很有效：对QKV投影生成的矩阵做FFT时，把变换结果缓存下来。因为attention计算往往会重复很多次，同一组输入在不同层里会反复参与计算，缓存频域表示能省下大量重复变换的时间。我把这个缓存做到显存池里，实际运行中整个attention模块的速度大约提升了35%到50%。
4.3 什么时候Attention值得上FFT
不是所有attention都要走FFT路线，它对场景是有要求的。
序列长度在512以下的时候，标准attention的 (O(L^2)) 计算也就是几十万次乘加，GPU处理这类规模本身就很快，FFT的变换开销相对就太高了。这种情况下不如老老实实优化数据布局和分块策略。
当序列长度超过2048，而且模型对长距离依赖要求不那么精细时，FFT方案的优势就开始出现了。再叠加qkv权重复用、批量推理需求这些条件，FFT方案在attention里的收益基本就实打实摆在那里了。
还有一点要提醒，FFT替换attention后，模型的训练行为和推理行为会发生变化。如果模型是已经预训练好的，直接换掉attention层可能会导致严重的精度恶化。这种情况正确的姿势是让上层也参与适配，不光是推理时替换一个算子就完事。
5. Softmax：最不适合FFT的算子，但有替代思路
softmax在这四个算子里，直接和老老实实说结论：FFT对它基本没用。
5.1 为什么Softmax用不上FFT
Softmax的数学形式是逐元素指数归一化：
[ \text{softmax}(x_i) = \frac{e^{x_i - \max(x)}}{\sum_j e^{x_j - \max(x)}} ]
这是一个典型的内存带宽密集操作，而非计算密集。它的瓶颈不在乘加运算，而在读取数据和输出结果的速度。FFT能加速的核心是"计算密集但结构性重复"的操作，对纯数据搬运操作一点忙都帮不上。
我把一个大型输入做softmax时，把FFT方案和直接指数方案都测过。FFT方案不但没变快，反而因为额外的变换开销变慢了3倍不止。这个结果其实没什么好惊讶的，就是把软盘的传动轴装到硬盘上，转不起来。
5.2 Softmax真正的优化方式是数值稳定和数据复用
在实际工程里，优化softmax的方向应该放在这些地方：
- 用在线归一化方法，也就是一次遍历同时计算最大值和求和，避免二次访存；
- 指数运算用近似指令，或者查表，但要注意精度对齐；
- 在attention结构里把softmax和前面的QK^T以及后面的PV融合到一起，减少kernel launch和中间张量的读写。
这些优化的共同特点都是围绕访存和融合做文章，跟FFT一点关系都没有。
5.3 Softmax在FFT加速attention中的角色
不过诚实地讲，在我之前做attention的FFT优化方案里，softmax确实不是完全不能碰，因为它出现在attention内部。
我在做频域attention近似时，把softmax这一步也顺带做了重排。做法是这样的：在频域完成token混合之后，不再显式构造完整的注意力矩阵，也不对那个矩阵做完整softmax，而是用一个可学习的门控缩放来替代softmax的归一化作用。这样做损失了一些理论上的严谨性，但实验上把整体精度控制在可接受范围内，同时省掉了矩阵级softmax的庞大访存开销。
这个经验只适用于我这种低频近似的attention架构，如果你在标准attention里试图绕过softmax，我劝你别冒险——那会带来不可控的数值问题。
5.4 给Softmax场景下的工程建议
如果遇到"怎么优化softmax"这类问题，正确反应不是去查FFT怎么用，而是检查你的数据流里softmax是不是被反复计算了太多次、中间结果是不是没有被缓存、kernel边界是不是把访存效率拖垮了。
把这三件事做明白，softmax的性能就能提升一大截。FFT是无辜的，也是无用的。
6. 实测对比：四个算子的FFT收益到底有多大
上面讲了不少原理和思路，这一节放一批我实测整理出来的数据。测试环境是某款主流GPU，批量大小为16，输入尺寸做了统一规模处理，精度为fp32。
6.1 测试配置和数据汇总
| 算子 | 原始耗时（ms） | FFT优化耗时（ms） | 加速比 | 结论 |
|---|---|---|---|---|
| CNN 5×5卷积 | 12.4 | 6.8 | 1.82x | 值得优化 |
| CNN 3×3卷积 | 8.1 | 10.5 | 0.77x | 不值得优化 |
| GEMM（循环结构） | 7.9 | 4.3 | 1.84x | 有条件值得 |
| GEMM（稠密矩阵） | 10.2 | 11.7 | 0.87x | 毫无价值 |
| Attention注意力 | 23.5 | 14.2 | 1.65x | 有条件值得 |
| 序列内softmax | 3.8 | 12.1 | 0.31x | 完全错误方向 |
这组数据基本验证了我前面所有的判断：FFT对卷积和注意力有效，对带循环结构的GEMM有条件有效，对稠密GEMM和普通softmax就是纯负优化。
6.2 一条数据背后的工程含义
加速比看着都是1.8倍上下，但落地难度差别巨大。CNN的FFT优化路径清晰，依赖现成的FFT库，难点主要在内存布局和批次变换的组织。attention的FFT优化本质上是算法替换，要动模型结构，工程复杂度和风险都高了一截。带循环结构的GEMM遇到的概率本身就不高，能匹配到结构矩阵的机会少之又少。所以如果要把这几个算子的FFT套路全用上，投入产出比差异很大。
我的排序是：CNN >= 带结构的GEMM > attention > softmax。
6.3 性能测试中的隐藏陷阱
最后吐槽一下测试陷阱。FFT性能对变换尺寸极其敏感，尺寸是2的幂和不是2的幂，性能差距可能达到4到5倍。我这里用到的尺寸都做了对齐处理，你要是照抄这组数据，务必先把尺寸对齐方法搞一致。
还有一个点是批次大小。FFT对批次大小的敏感度比普通kernel更高，小批次下变换开销占比大，优势不明显；批次大了之后，批量变换的并行优势才能发挥出来。测下来批次从4涨到16的收益，比从16涨到64更明显，批次大于64后基本就到天花板了。
7. 工程落地的几条经验总结
最后把这段摸索过程里最有价值的经验整理成几条，供后来者直接参考。
第一条：先问算子的数学结构，再谈FFT。FFT能加速的操作必须满足卷积定理的适用条件。逻辑上要是乘法、加法、分块共享，能转成乘积或卷积结构才有戏。普通矩阵乘、softmax这类直接硬上，一定亏。
第二条：别省预处理的复杂度。FFT方案里，输入数据的变换、卷积核的变换、批量变换的组织，都是在真正计算乘法之前的额外开销。这段开销如果没处理干净，后半段所有的复杂计算都白搭。我当时花了不少精力在批量FFT的规划上，最后收益就体现在这里。
第三条：精度不能最后测。FP32下问题不大，但混合精度或者低精度推理下，FFT的频域累积误差会成为模型精度劣化的暗雷。设计一开始就要确定FFT路径的精度档位，尽量隔离在fp32区域执行复数运算，不要想着最后统一校准。
第四条：FFT不一定只做FFT。像attention这种场景，我最后落地的时候，核心不是暴力FFT变换，而是用FFT做低频近似加上分块attention融合。算法替换和库函数优化是两条相辅相成的路，别被"用FFT库"这个字面意思框死了。
第五条：模型如果是从头训练，把FFT结构设计进去。推理阶段的FFT替换往往受限于训练时模型没有适应频域近似带来的偏差。如果能在训练阶段就让模型过一遍频域结构，最终精度和推理速度的配合会比事后替换好太多。能走这条路的时候，一定优先走这条路。
说点私心话。FFT优化算子这事，看起来是数学结构的问题，实际上更多是工程取舍的艺术。每个算子的数据规模、计算模式、存储访问特性都不一样，有没有一次碰撞到那个核心矛盾，决定了FFT到底是"加速器"还是"绊脚石"。我在这几个算子上跑完一圈，最大的收获反而是更清楚地认识到了每个算子本身的边界。你在选型的时候，也别忘了先量好自己的数据节奏和硬件脾气，再去谈用不用FFT这件事。
