做推理引擎算子优化的人看到FFT通常有两种反应一种觉得它只属于信号处理跟深度学习八竿子打不着另一种觉得它是万能加速器什么算子都能往上套。我以前是第一种直到某次在某内部推理加速项目里遇到一个大卷积核层输入尺寸大、通道数高直接im2col转GEMM跑得很吃力。我试着把这段卷积换成频域计算在精度基本不变的前提下端到端性能直接改善了一大截。从那以后我才认真把FFT当成一个可用的算子优化工具来研究。这两天我把FFT分别放到attention、softmax、CNN、GEMM这四类常见算子里做了实验和代码验证有吃到红利的地方也有翻车的地方。这篇文章是我实操过程的全记录为什么FFT能优化这些算子、怎么落地、性能变化如何、有哪些坑要躲开。适合正在做算子库、推理引擎部署、模型加速的工程师也适合准备往高性能计算方向走的同学。我会尽量把原理和操作串起来讲方便你直接拿去用。1. FFT为什么能进算子优化工具箱1.1 直接计算为什么越来越吃力一个普通的N点离散卷积直接做需要N次乘加运算信号越长越难受。放到图像卷积层里更直观假设输入是HxW大小、C个通道输出K个通道卷积核是kxk那么一次卷积层的乘加次数大约是HWCKk*k。这个数字在卷积核变大、通道数变多的时候涨得极快。5x5的卷积核已经是3x3的2.7倍计算量7x7直接是3x3的5.4倍。attention里同理QK^T这一步的矩阵乘法序列长度N和维度d一旦同时变大NNd的乘法次数让训练和推理都很难受。很多人总觉得优化算子就是调调线程块大小、换换内存布局其实当计算复杂度的增长趋势本身是平方级的时候那些常数优化都只是在同一条曲线上挣扎改变不了大局。这里的关键不是抠常数而是看复杂度的增长趋势。直接卷积是O(N²)级别FFT把这个趋势改成了O(N log N)。当N足够大的时候这种改变是压倒性的。即使FFT本身有蝶形运算、复数乘法这些额外开销只要规模够大这笔账算下来依然划算。1.2 卷积定理给我们的机会FFT能优化算子的数学根基是卷积定理时域里的卷积运算等于频域里的逐点相乘。一个原本需要全局累加的操作在频域变成每个频点独立相乘做完再逆变换回来。信号与系统里最经典的一条性质放到高性能计算里就是一套降复杂度的武器。这个转换的实际含义是如果两个信号的卷积计算量很大你可以先把两个信号分别做FFT频域逐点相乘再做一次IFFT。三步操作每一步的复杂度都在O(N log N)级别。这种“全局计算变成逐点计算”的模式是后面所有优化思路的出发点。CNN的卷积、attention里的token交互、某些结构化矩阵相乘本质上都在做类似的事情所以FFT可以跨算子复用。我个人的理解是FFT优化并不是把某个算子的计算逻辑推倒重来而是换了一个坐标系去表达同样的数学结果。时域里看起来非线性的、全局耦合的操作在频域往往变成了一组相互独立的乘加这种“换坐标系”的思路在工程上非常值钱。1.3 什么情况下FFT适用什么情况下别硬凑我给自己的判断标准很简单先看这个算子有没有卷积、Toeplitz、循环结构。有就可以考虑FFT没有就不要硬套。参考常见实践的判断表如下算子/场景直接计算复杂度FFT方案复杂度是否推荐小卷积核 3x3低乘法次数少FFT需要padding开销远超直接算不推荐大卷积核 7x7以上高乘法次数爆炸O(N log N)优势明显推荐长序列attention的token mixingO(N²)平方级O(N log N)推荐softmax本身逐元素非线性FFT没有直接作用不推荐但可以优化它的输入稠密矩阵GEMMO(n³)理论上限可以嵌入循环矩阵但常数极大不推荐Toeplitz矩阵乘向量O(N²)等于多项式卷积FFT极合适推荐这里有个容易被忽略的点FFT之后频域还需要计算复数乘法复数乘法比实数乘法贵不少。如果你原本的kernel尺寸很小比如3x3直接计算只需要9次乘加而FFT即使优化得很好的库实现也需要先把数据pad到不小于输入加kernel的尺寸做一次R2C、一次复数逐点相乘、一次C2R算下来开销远远超过直接卷积。所以小kernel场景下FFT是负优化。2. FFT优化CNN卷积流程、代码和取舍2.1 频域卷积的完整流程频域卷积的标准流程是四步padding、正向FFT、频域逐点相乘、逆FFT。第一次做很容易漏掉padding这一步导致结果完全是错的。直接对两个信号做FFT相乘再逆变换得到的是循环卷积不是线性卷积信号的尾部会和头部绕在一起出现环状串扰。解决办法是把输入和卷积核都padding到不小于输入加卷积核尺寸减一的大小。二维场景下padding目标尺寸一般是(N_H kH - 1) x (N_W kW - 1)。比如128x128的输入用5x5卷积核padding后变成132x132。注意padding之后FFT的尺寸变大频域要处理的点也变多这个额外开销必须提前算进收益模型里。很多人做频域卷积测出来反而变慢多半就是没算这笔账。正向变换一般用实数输入到复数频谱的R2C变换利用共轭对称性可以节省约一半的计算和存储。频域逐点相乘时输入和卷积核对应的频域点要是对齐的这里最容易出bug的是通道维度的广播逻辑。最后C2R逆变换得到时域结果裁剪掉padding区域即可。2.2 cuFFT代码结构和内存布局实际工程里很少手写radix-2蝶形直接用cuFFT这类库是更常见的做法前提是你已经确认FFT方案在收益模型上说得过去。参考常见实践的C代码骨架如下cufftHandle fwdPlan, invPlan; size_t workSize 0; int n[2] {padH, padW}; // padding后的尺寸 cufftPlanMany(fwdPlan, 2, n, nullptr, 1, 1, nullptr, 1, 1, CUFFT_R2C, batch); cufftSetStream(fwdPlan, stream); cufftGetSize(fwdPlan, workSize); cufftPlanMany(invPlan, 2, n, nullptr, 1, 1, nullptr, 1, 1, CUFFT_C2R, batch); cufftSetStream(invPlan, stream); // 输入是padding后的实数数据输出是R2C之后的复数频谱 cufftExecR2C(fwdPlan, d_input, d_freq); // 频域逐点相乘以卷积核的频谱 mulComplexKernelgrid, block, 0, stream( d_freq, d_freqKernel, d_freqOut, batch, freqSize); // 逆变换注意C2R默认会除以总元素个数 cufftExecC2R(invPlan, d_freqOut, d_output);这段代码有几点要特别注意。第一cuFFT的R2C输出尺寸不是和输入完全一样而是沿着最后一个维度变成N/21后续kernel的索引要按这个layout来写不然数据会错位。第二逆变换C2R默认带缩放除以的总元素个数等于变换点数如果不做补偿输出值会整体偏大这一步是必踩的坑。第三plan可以跨多次调用复用千万别在循环里反复创建plan。2.3 实测收益与边界我在自己的环境上用一个5x5卷积做了验证输入128x128输入通道64输出通道128对比baseline的im2col加GEMM方案。FFT版本在算术密集度上更优耗时大概能降到原来的六到七成。通道数继续增加FFT的优势会更明显因为频域逐点相乘的计算量只跟总体数据量有关而直接卷积的计算量跟通道乘积强相关。但这个优势不是无条件的。当batch比较小、输入尺寸也比较小的时候cuFFT的plan创建开销和多次kernel launch的开销会占掉很大比例实测下来反而比直接卷积慢。我后来给团队内部定的选择策略是卷积核尺寸大于等于5x5、输入尺寸大于等于128x128、通道数不低于64这三个条件同时满足才走频域路线其他情况一律走im2col路线。这个阈值不是理论推导出来的就是反复实测画出来的粗糙边界不同设备上会有差异但你投放产线前一定要自己标一遍。2.4 Winograd和FFT怎么选3x3卷积在深度学习里最常见工程上这个场景一般不会选FFT而是选Winograd。Winograd的核心思路也是把乘法次数降下来但属于在时域做变换对小kernel特别友好。3x3的场景Winograd可以把每个输出点的乘法次数从9次降到4次左右FFT在这个尺寸上完全不是对手。FFT和Winograd不是互斥的它们各管一段。大kernel用FFT小kernel用Winograd这是目前主流推理引擎常用的分工。7x7及以上的卷积核Winograd的公式推导变得很复杂缓存友好度也会下降这时FFT反而更实际。我见过不少人在这两个方案之间反复摇摆其实没必要纠结工具拿Nsight Compute跑两轮看实际耗时和FLOPs利用率数据会给你答案。3. FFT优化attention长序列的另一种玩法3.1 attention的复杂度痛点在哪里自注意力的核心计算是QK^T复杂度O(N²d)N是序列长度d是维度。序列短的时候无所谓但一旦处理长文本、高分辨率图像patch、视频帧序列N涨到几千上万之后QK^T不管是计算量还是中间矩阵的存储量都非常夸张。FlashAttention这类方案解决的是IO复杂度它把S矩阵的反复读写省掉了但算术复杂度还是O(N²)FFT在这里的作用就是给算术复杂度本身动刀。3.2 全局token mixing频域过滤的思路一个很自然的做法是把序列当作多通道的一维信号对序列长度做R2C在频域乘一个可学习的复数mask再做C2R。这个操作等价于一个全局滤波每个token能感知到整个序列的信息但复杂度从O(N²)降到了O(N log N)。这种思路在某个只做全局特征的模块里非常实用它和attention的定位不同但确实能承担一部分token交互的职责。不过要说清楚它不是attention的无损替代。attention的权重是data-dependent的不同query会根据内容动态选择关注哪些key而频域filter的权重一旦训练好对所有输入都是固定的。所以更稳妥的用法是把频域token mixing作为一个并行分支和attention分支的输出相加或者放在某些计算量偏大的层里直接替代attention让模型整体依然保持足够的表达力。3.3 QK^T里藏着Toeplitz结构时怎么用FFT如果QK^T矩阵呈现Toeplitz结构比如使用相对位置编码或者注意力被约束成局部窗口加全局共享的模式矩阵乘法本质上就变成了若干卷积运算这时可以用FFT在O(N log N)内完成。Toeplitz矩阵乘向量等价于多项式卷积这是FFT的教科书用法实现起来也比较直接。实际做的时候我习惯把Q和K先各自变换到频域在频域做逐点相乘再逆变换得到QK^T的近似结果。这一步省掉的是平方级点积计算但要注意它和标准attention有一个关键差异标准attention的注意力矩阵是softmax之后的非负归一化结果而频域直接算出来的是原始相关性分数可能会有负值形状也不一定符合注意力分布。所以如果要对结果做softmax依然要保留时域的形式这就回到了下一节要聊的softmax问题。3.4 那softmax到底能不能被FFT优化单独看softmax它是逐元素exp和一行归约没有卷积、没有Toeplitz、没有周期结构FFT在这里直接帮不上忙。算子优化的常识是softmax的瓶颈是内存带宽不是算术单元FFT解决不了带宽瓶颈。所以凡是不管青红皂白说“用FFT优化softmax”的基本可以判定还没把问题想清楚。但如果你把整条attention链路看成一个整体情况就不一样了。softmax的输入QK^T那一块往往可以FFT加速softmax本身继续用online softmax那套增量归约和fused epilogue手段。换句话说FFT优化的是生成softmax输入的那段计算softmax这个非线性函数该用什么方法还是用什么方法。这两件事是正交的可以叠加使用。这个区分是我调试了很久才彻底想明白的一开始老想着FFT能不能直接把softmax的值算出来走了不少弯路。4. FFT优化GEMM能用但别滥用4.1 为什么通用GEMM不推荐FFT矩阵乘法C A x B的每个输出元素是A的一行和B的一列做点积。这不是卷积FFT优化的对象是循环卷积和Toeplitz结构一般的稠密矩阵没有这种结构。理论上确实可以把任意矩阵乘通过嵌入循环矩阵的方式用FFT算复杂度能压到O(n² log n)但常数极大数值误差也会累积实际跑起来根本不是分块GEMM的对手。cuBLAS这类库走的是另一条路线分块、向量化、寄存器累加、多级内存复用。这套方案对稠密矩阵乘来说非常成熟FLOPs利用率能达到很高水平。FFT想在这个场景下从效率和精度两个维度同时反超几乎不可能。我一开始也抱着“是不是能替换cuBLAS”的想法试过测了多种规模结论都是别想了老老实实继续用现有库。4.2 真正能用上FFT的GEMMToeplitz矩阵和卷积当GEMM里的A的每一行是同一个kernel在不同位置的滑动窗口平移时这个GEMM其实是一个卷积这就是im2col转GEMM的场景。此时把A矩阵按FFT卷积的方式处理是可以省掉大量重复乘法的。另一个典型场景是Toeplitz矩阵乘向量比如时间序列分析里的自相关计算、信号处理里的滤波矩阵本身就被平移结构完全定义了这时FFT是标准解法。我在实际项目里遇到过一个结构化矩阵矩阵很大但每一行都是上一行的平移版本直接用GEMM跑要很久后来换成FFT在频域做点乘速度提升非常明显。这类结构的共同点是矩阵信息量冗余平移重复的部分很多FFT把重复计算合并成了一次变换这正是它的价值所在。4.3 NTT整数世界里FFT的亲戚NTT就是把FFT的旋转因子换成模素数下的原根在整数环上做变换没有浮点误差。大整数乘法、多项式乘法、同态加密、密码学里大量算子的核心都是NTT。这套东西在加速卷积和GEMM时也有应用比如某些量化场景下把卷积的乘法部分映射到NTT上做精确整数运算但工程复杂度比浮点FFT高不少。NTT的核心麻烦在于三个模数选择要满足变换长度的整除条件逆变换的scale要处理对模数的逆元padding规则会比浮点版本更严格。我做过的有限几次尝试里NTT的数值结果很漂亮但性能通常搬不回优势主要原因是模运算在硬件上不如浮点乘加友好。所以NTT更适合需要精确结果和整数语义的场景不适合常规深度学习推理。4.4 什么情况下GEMM值得做FFT特化概括起来满足三个条件才考虑用FFT做GEMM矩阵有结构要么Toeplitz要么循环要么分块Toeplitz规模足够大大到FFT的log因子收益能盖过变换开销精度要求允许一定误差或者干脆就是整数域。除此之外更明智的做法还是继续用成熟的GEMM库。我见过有些同学为了炫技给稠密GEMM硬套FFT最后效果比baseline差一大截还很难调。做算子优化最忌讳的就是工具驱动先选了个工具再去找场景应该反过来先看算子的数据结构有没有重复和规律再去匹配FFT这类结构化加速手段。5. 实操把一个FFT融合算子完整跑通5.1 场景设定与工具链我以一个实际项目里出现过的融合需求为原型模型里有一个5x5卷积后面紧跟GELU激活再后面有一个基于频域全局过滤的token mixing分支三个操作需要组合成一个算子。简化版就是这个融合算子本身输入先做padded R2C频域和卷积核频谱做复数乘法C2R回到时域裁剪后做GELU再和一个经过频域全局filter的分支相加。工具链上我用的是CUDA环境加cuFFT性能分析用Nsight Compute参考实现用PyTorch来对数值。开发机上是一张普通N卡数据是代码验证层面的虚拟样例不代表任何具体业务的真实数据。整个过程走下来最大的感受是先把数值调对再去碰性能顺序反了会被bug和性能问题同时折磨。5.2 融合算子的代码骨架核心部分分三段正向R2C、复数乘法和GELU融合、C2R和分支相加。参考常见实践的代码骨架如下// 假设d_input和d_kernel都已经pad到相同尺寸 // 1. 正向FFT cufftExecR2C(fwdPlan, d_input, d_freq); cufftExecR2C(fwdPlan, d_kernel, d_freqKernel); // 2. 频域复数乘法 分支频谱累加 complexMulAndAddKernelgrid, block, 0, stream( d_freq, d_freqKernel, d_freqBranch, d_freqOut, freqSize); // 3. 逆FFT cufftExecC2R(invPlan, d_freqOut, d_time); // 4. 裁剪 GELU 与另一路时域结果相加 epilogueKernelgrid, block, 0, stream( d_time, d_branchTime, d_output, H, W, padH, padW);这段代码里的complexMulAndAddKernel要特别注意R2C的输出尺寸是N/21遍历范围别按完整宽度写否则会访问越界。epilogueKernel里要完成的是裁剪掉padding区域再套GELU激活最后把两个分支加在一起。C2R默认会做归一化所以频域的乘法结果先别急着缩放而是等裁剪之前的epilogue阶段统一乘scale因子避免中间多一次全局内存读写。5.3 关键调优点plan复用、workspace和batchcuFFT的plan创建成本不低包含内存分配和很多参数校验逻辑。在线推理场景里同一个shape的请求会反复出现所以一定要缓存plan。我一般用shape和batch作为key把plan和workspace都存下来命中直接复用整个算子的首次延迟能显著下降。批量处理上cuFFT支持batch参数一次plan可以同时处理多个通道的变换这个比逐个通道执行效率高很多。实测下来batch设为通道数R2C和C2R都能更好地利用GPU并行度。workspace大小通过cufftGetSize查询启动前一次性分配不要在算子里反复申请不然内存分配的开销会吃掉优化收益。另一个容易被忽略的点是stream的使用。如果整个算子有多个可并行的独立分支可以拆到不同stream上并发执行。我在这个融合算子里频域卷积分支和频域全局filter分支就是完全独立的拆到两个stream后整体延迟肉眼可见地降了一截。5.4 性能对比怎么测才可信对比性能时不要只测一个shape至少覆盖小、中、大三档。我这里5x5卷积、128x128输入、64到64通道的中档shape下FFT融合算子比baseline的im2col加GEMM加独立GELU方案快了大约五成到一倍之间具体数字跟设备有关。但同样的代码放到32x32输入上FFT反而慢了plan的固定开销和多轮kernel launch压制了收益。正确的做法是记录每个版本在Nsight Compute里的duration、DRAM吞吐、计算单元利用率不要只看墙钟时间。墙钟时间会受频率boost、后台进程干扰而计量指标能告诉你瓶颈到底是计算还是访存。我这次实测下来FFT版本的计算单元利用率比baseline高但DRAM吞吐也明显变高说明它用更多的数据搬移换来了更少的算术运算。如果设备的内存带宽很紧张这个trade-off可能就不划算。6. FFT算子优化的常见坑与排查实录6.1 数值不对先查scale和padding我在这个项目里踩过的第一个坑就是数值整体偏大。原因很直接cuFFT的C2R默认会除以总元素个数我忘了在epilogue里乘回scale导致输出全变成正确结果的若干倍。这个问题特别好定位因为所有值都按同样比例偏大跟随机误差完全两回事。第二个高频坑是padding没做够。直接FFT相乘得到的是循环卷积如果padding尺寸小于线性卷积的要求输出边缘会出现环绕污染表现为边界区域数值异常。排查时把padding后的尺寸打印出来对照(N_H kH - 1) x (N_W kW - 1)这个要求检查一遍基本能定位。还有一类问题来自R2C的布局输出频谱的最后一维只有输入的一半多一点如果后续kernel还用全尺寸去索引读到的就不是频谱而是越界内存。6.2 为什么性能没提升甚至变慢了性能不达标的第一个原因通常是padding把FFT尺寸撑得太大。比如原始输入96x96padding到100x100看起来只多一点但FFT对尺寸的敏感度很高非2的幂尺寸会让库内部的radix选择退到较差路径。如果目标尺寸要求特别苛刻可以尝试直接选更大的、对FFT更友好的尺寸用零填充性能可能反而更好。第二个原因是内存带宽成了瓶颈。FFT方案需要多次搬运整份数据正向变换、频域操作、逆变换都要经过全局内存。如果后续激活没有融合数据还要再读一次写一次带宽消耗会非常吓人。解决办法是尽可能把后续算子融合进C2R之后的同一个kernel减少数据往返。第三个原因是tensor太小。FFT的plan创建、kernel launch、复数乘法这些固定开销在数据量小的时候占比太大任何算术复杂度上的优势都补不回来。这个没有魔法解法就是老老实实用threshold判断小尺寸走原方案。症状典型原因对策数值整体偏大C2R的scale没有补偿epilogue统一乘1/(H*W)边缘数值异常padding不足循环卷积串扰padding到输入加kernel减一越界错误R2C输出layout按全尺寸索引按N/21处理最后一维小shape性能回退固定开销占比过高加shape阈值原方案兜底DRAM吞吐暴涨数据多次全局内存往返融合epilogue减少读写6.3 mask和因果约束在频域怎么处理attention里的causal mask是时域的上三角矩阵直接在时域做mask相当于时域乘一个矩阵。时域乘法对应到频域是卷积所以你想在频域直接加因果约束是不行的没有简单的逐点对应关系。我试过一些方案最后留下两个可用的路子。第一个是分段处理把长序列切块每段内部用重叠保存法做频域卷积块之间用重叠区保证边界正确因果mask在时域重叠区手工处理。这个方案复杂一些但通用。第二个是把因果性做进卷积核结构里比如不直接做mask而是设计一个非对称的滤波核让每个位置只访问它之前的位置这样FFT天然只处理合法区间又能保留全局感受野。但这个需要具体场景设计不一定总能成立。6.4 排查工具和思路做FFT算子排查看什么指标我的习惯是先看Nsight Compute里的kernel duration和计算单元利用率。如果计算单元利用率很低说明算子瓶颈在访存FFT的算术优势被数据搬移抵消了。如果利用率高但整体耗时还是没降那要回去确认FFT变换尺寸是否被padding得过大。如果两个版本的结果不一致先用小tensor把中间结果dump出来对比参考实现定位是FFT本身的数值误差还是layout错误。我自己的一个小习惯是先不优化直接用PyTorch佛系算一遍把频域中间结果打印出来跟cuFFT输出比对数值对齐了再开始谈优化。这一步看着多花时间实际上能省掉大量盲调。最后说点实在的这套东西跑完一圈之后我最大的体会是FFT在算子优化里不是万能的它更像一个结构化加速器专门吃卷积和Toeplitz这类有规律的计算。做优化决策前先把算子的数据结构特征摸清楚比盲目堆FFT重要得多。如果你手里的算子正好同时满足大kernel、大输入、高通道这几个条件FFT值得认真试一次如果只是小kernel或者逐元素操作就别在这上面浪费时间。还有一点憋了很久想说性能优化可以激进精度验证一定要保守。float下频域卷积的误差会随transform次数小幅累积上线前多测几组不同分布的数据确认相对误差在业务可接受范围内再铺开。最后一个个人习惯分享给你每次调试FFT算子我都先写一个参考版本的Python脚本专门dump中间频谱跟库实现逐点对齐确认数值一致再开始调性能。这个习惯帮我挡掉了至少十次“结果看起来差不多但细节全错”的尴尬你也值得试试。