1. 从一次性能瓶颈说起为什么FFT能成为通用算子优化的钥匙做过深度学习推理优化的人大概都有过这种体验模型结构没怎么变算力也没少堆但端到端延迟就是卡在某个数字下不去。你打开profiler一看attention、softmax、卷积、矩阵乘各占一块单独看每个kernel都不算慢但拼在一起就是达不到预期。这时候很多人第一反应是去调batch size、换更快的库、或者干脆上更激进的量化但往往收效有限。我最近在做一个推理加速的项目时重新把FFT快速傅里叶变换这个老古董翻了出来。原因很简单卷积本质上就是频域乘法这个结论在信号处理领域是常识但在深度学习算子优化里很多人只把它当成一个教科书里的等价变换没真正当成工程手段来用。而一旦你认真把FFT当成一个通用的算子加速工具会发现它能撬动的远不止卷积——attention里的相关运算、softmax的归一化、甚至GEMM在特定形状下的分块策略都能借上力。这篇文章想聊的就是这件事怎么用FFT的思路去优化其他kernel和算子包括attn、softmax、cnn、gemm这几类。不是纯理论推导而是从工程落地角度讲清楚哪些场景值得用、怎么用、坑在哪、实测效果如何。适合有一定CUDA或算子开发基础、正在做推理性能优化的同学参考也适合对频域方法感兴趣但没实际动过手的人。先说结论FFT不是万能药它有自己的甜点区。用对了某些kernel能拿到2到5倍的加速用错了反而比直接算还慢。关键是要理解它的适用边界以及怎么和现有算子融合策略配合。2. 核心思路拆解FFT为什么能加速这些算子2.1 卷积定理一切加速的起点卷积定理说的是时域或空域的卷积等价于频域的逐点相乘。用公式表达就是F(f * g) F(f) · F(g)其中F表示傅里叶变换*表示卷积·表示逐点乘法。反过来如果你想在频域做乘法再变回来就得到f * g F⁻¹(F(f) · F(g))这个等式看起来简单但工程含义很重。直接做空域卷积复杂度是O(N²)对于一维或者O(N²·K²)对于二维卷积核大小K。而通过FFT复杂度降到O(N log N)。当卷积核较大或者输入尺寸较大时这个差距非常可观。但这里有个关键点很多人忽略FFT加速卷积的前提是卷积核足够大。如果卷积核只有3x3直接算反而更快因为FFT的常数因子和变换开销摆在那里。一般来说当卷积核尺寸超过7x7或者等效感受野很大时FFT方法才开始有优势。这是第一个要记住的边界条件。2.2 从卷积扩展到attention和softmaxattention的核心计算是QK^T这本质上是一个矩阵乘法但在某些变体里比如linear attention或者某些kernelized attention它等价于一个相关运算。相关运算和卷积是近亲只差一个翻转。所以当attention的序列长度很大、head dimension相对固定时用FFT做相关运算可以降低复杂度。softmax看起来和FFT没关系但如果你仔细看softmax的归一化过程它涉及对指数化后的值做求和。在某些近似方法里比如基于随机特征的方法这个求和可以转化为频域上的积分。更实际的是当softmax和卷积组合在一起时比如在卷积注意力模块里可以把整个链路放到频域里做减少中间结果的往返。GEMM的情况更微妙。标准GEMM是O(N³)的FFT本身不能直接降低矩阵乘的复杂度。但在特定场景下比如当矩阵具有Toeplitz结构或者循环结构时矩阵乘可以转化为卷积进而用FFT加速。这在某些权重共享或者结构化剪枝的场景里会出现。2.3 方案选型的核心考量决定要不要用FFT我一般看三个维度维度适合FFT不适合FFT卷积核/等效窗口大小大于7x73x3或更小输入尺寸较大如224x224以上很小如32x32以下精度要求允许一定数值误差需要bit级精确硬件有良好FFT库支持FFT实现效率低数值误差这点要特别说明。FFT涉及浮点运算的多次累加相比直接卷积误差会略大。在fp32下通常没问题但在fp16或者int8下要小心。我实测过fp16下FFT卷积的误差大约在1e-3量级对大多数推理任务够用但对某些敏感任务可能需要fp32。3. 核心细节解析与实操要点3.1 FFT卷积的实现细节先讲最成熟的FFT卷积。假设你有一个输入特征图X尺寸H×W和一个卷积核K尺寸Kh×Kw标准做法是把X和K都padding到(HKh-1)×(WKw-1)避免循环卷积的混叠对两者分别做2D FFT频域逐点相乘做2D IFFT裁剪到有效区域这里有几个实操要点Padding策略必须padding到至少(HKh-1)×(WKw-1)否则会出现循环卷积的wrap-around错误。我见过有人为了省内存只padding到H×W结果边缘出现明显伪影。另外如果尺寸不是2的幂FFT库通常会自动处理但效率会下降。可以手动padding到最近的2的幂或者5-smooth数如2^a·3^b·5^c这样FFT最快。批处理维度实际推理时batch size通常大于1可以把batch维度合并到FFT的批处理里。cuFFT支持batch FFT一次调用处理多个样本比循环调用快很多。内存布局FFT对内存访问模式敏感。如果输入是NCHW布局做2D FFT时最好把H和W维度连续存放。我一般会先做一次permute或者用stride技巧让FFT的输入是连续内存。3.2 Attention中的FFT应用Attention的FFT优化分两种情况。第一种是标准attention的近似。当序列长度L很大时QK^T的复杂度是O(L²·d)其中d是head dimension。如果d远小于L可以考虑把Q和K的每一行看作一个信号做一维FFT然后频域相乘再IFFT。这样复杂度降到O(L·d·log L)。但要注意这只在d较小比如d64时有优势因为FFT的常数因子和d有关。第二种是linear attention或者kernelized attention。这类attention把softmax(QK^T)近似为φ(Q)·φ(K)^T其中φ是某个特征映射。如果φ选择得当比如基于随机傅里叶特征整个计算可以转化为频域上的操作。我试过一个基于FFT的linear attention实现在L4096、d64时比标准attention快约3倍但精度损失在1%左右。实操上attention的FFT优化要注意数值稳定性频域相乘后再IFFT结果可能超出fp16范围建议中间结果用fp32。mask处理attention通常有causal mask频域方法处理mask比较麻烦。一种做法是先把mask应用到频域但这会破坏卷积定理的前提。更实际的做法是分段处理把序列切成块块内用FFT块间用标准方法。head维度如果head dimension不是2的幂padding到最近的2的幂会引入额外计算。可以尝试混合策略对部分head用FFT部分用标准方法。3.3 Softmax的频域近似Softmax本身是逐元素的FFT不能直接加速。但在以下场景可以借力场景一softmax 卷积。如果softmax的输出紧接着一个卷积层可以把卷积放到频域做softmax的结果作为频域乘法的输入。这样省去了中间结果的往返。场景二大类别softmax。当类别数极大比如语言模型里的vocab size时softmax的归一化求和很耗时。有一种近似方法是用FFT做快速求和但精度损失较大我一般不建议在最终输出层用。场景三softmax的梯度计算。反向传播时softmax的Jacobian是稠密的但如果用频域方法近似可以降低计算量。这在训练场景有意义推理场景一般不需要。我实测下来softmax单独用FFT的收益很小除非它和卷积或attention组合在一起。所以实操建议是不要为了FFT而FFT要看整个计算图。3.4 CNN中的FFT卷积实战CNN是FFT卷积最成熟的应用场景。我以一个实际的ResNet变体为例讲一下怎么落地。假设有一个卷积层输入是256×56×56卷积核是7×7输出256通道。标准卷积的计算量大约是256×256×56×56×7×7 ≈ 9.6 GFLOPs。用FFT的话把输入和卷积核padding到64×64567-162取最近的2的幂64对每个输入通道和输出通道对做FFT频域相乘并累加IFFT得到输出实测在V100上这个层用FFT比cuDNN的标准卷积快约1.8倍。但如果卷积核换成3×3FFT反而慢20%。所以7×7是FFT卷积的一个甜点。注意事项通道数FFT卷积对通道数敏感。通道数太少比如32FFT的批处理效率上不去。通道数太多比如512内存占用会很大。精度fp16下FFT卷积的误差累积比标准卷积明显建议关键层用fp32。融合FFT卷积后面如果接BN和ReLU可以把BN的参数吸收到频域乘法里ReLU则必须在IFFT之后做。3.5 GEMM的结构化加速GEMM用FFT加速的条件比较苛刻。只有当矩阵具有Toeplitz、循环或者近似低秩结构时才能转化为卷积。在实际模型里这种结构出现在权重共享的层比如某些轻量级模型里的深度可分离卷积其逐通道卷积可以看作Toeplitz矩阵乘。结构化剪枝后的层剪枝后权重矩阵可能呈现块循环结构。某些attention变体比如Toeplitz attention。我试过一个Toeplitz GEMM的FFT实现在矩阵尺寸1024×1024、秩为64时比标准cuBLAS快约2.5倍。但如果是稠密随机矩阵FFT方法完全没优势。实操建议先分析权重矩阵的结构再决定要不要用FFT。可以用奇异值分解或者结构检测工具先看一眼。4. 实操过程与核心环节实现4.1 环境准备与工具选型我用的技术栈是CUDA cuFFT cuBLAS配合PyTorch做验证。cuFFT是NVIDIA的FFT库成熟稳定支持batch和多种尺寸。如果要做自定义kernel可以用cuFFT的callback功能或者自己写FFT kernel不推荐除非有特殊需求。工具选型对比工具优势劣势适用场景cuFFT成熟、快、支持batch闭源、灵活性差通用FFT卷积cuBLASGEMM优化好不支持FFT标准GEMM自写kernel灵活开发成本高特殊结构PyTorch FFT易用性能一般原型验证我一般先用PyTorch FFT做原型验证精度和大致加速比再用cuFFT做工程实现。4.2 FFT卷积的完整实现步骤以一个7×7卷积为例完整步骤如下步骤一确定padding尺寸。输入56×56卷积核7×7输出50×50。FFT尺寸取64×64最近的2的幂。步骤二准备频域权重。把卷积核padding到64×64做2D FFT得到频域权重。这个可以预计算推理时直接用。步骤三输入FFT。把输入padding到64×64做2D FFT。如果是batch输入用cuFFT的batch模式。步骤四频域相乘。对每个输入通道和输出通道对做逐点复数乘法然后累加。这一步是计算密集的可以用cuBLAS的复数GEMM来加速。步骤五IFFT。对累加结果做2D IFFT取实部。步骤六裁剪和激活。裁剪到50×50然后做BN和ReLU。关键代码片段伪代码# 预计算频域权重 weight_freq fft2(pad(weight, (64, 64))) # 推理时 input_freq fft2(pad(input, (64, 64))) output_freq complex_gemm(input_freq, weight_freq) output ifft2(output_freq).real output crop(output, (50, 50)) output relu(bn(output))实测在V100上这个实现比cuDNN的7×7卷积快1.8倍。但如果用TensorRT的融合卷积差距会缩小到1.3倍左右。4.3 Attention的FFT实现以linear attention为例实现步骤步骤一特征映射。把Q和K通过φ映射到频域。φ可以选择基于随机傅里叶特征的方法。步骤二频域相乘。在频域做逐点乘法得到attention的频域表示。步骤三IFFT。变回时域得到attention输出。步骤四与V相乘。标准attention的后续步骤。实测在L4096、d64时比标准attention快约3倍。但精度损失约1%需要根据任务权衡。注意事项随机特征的数量φ的维度决定了近似精度。维度越高越准但计算量也越大。我一般取d的2到4倍。数值范围频域相乘后数值可能很大建议用fp32。maskcausal mask在频域不好处理建议分段。4.4 性能测试与对比我在V100上做了一组对比测试结果如下算子标准实现FFT实现加速比7×7卷积 (256×56×56)12.3 ms6.8 ms1.81x3×3卷积 (256×56×56)4.2 ms5.1 ms0.82xLinear attention (L4096, d64)8.7 ms2.9 ms3.0xToeplitz GEMM (1024×1024, rank64)3.5 ms1.4 ms2.5xSoftmax (vocab50000)1.2 ms1.3 ms0.92x从数据看FFT在7×7卷积、linear attention、Toeplitz GEMM上有明显优势在3×3卷积和softmax上反而慢。这印证了前面的边界条件。5. 常见问题与排查技巧实录5.1 数值误差问题问题FFT卷积在fp16下误差较大某些层输出出现明显偏差。排查先对比fp32和fp16的FFT结果看误差量级。如果误差在1e-3以上考虑关键层用fp32。解决混合精度策略——FFT用fp32逐点乘法和IFFT用fp16。或者用Kahan求和减少累加误差。5.2 内存占用问题问题FFT卷积的内存占用比标准卷积大很多尤其是大batch时。排查算一下频域张量的大小。对于64×64的复数FFT每个通道需要64×64×2×4字节32KB。256个通道就是8MBbatch32就是256MB。解决分块处理把batch或通道维度切分。或者用in-place FFT减少内存。5.3 尺寸不是2的幂问题输入尺寸不是2的幂时FFT效率下降。排查看cuFFT的plan创建时间。非2的幂尺寸的plan创建可能慢10倍。解决padding到最近的2的幂或5-smooth数。或者用混合基FFT。5.4 常见问题速查表问题可能原因解决方法输出有伪影padding不足padding到HKh-1精度差fp16误差累积关键层用fp32速度慢尺寸非2的幂padding到2的幂内存爆batch太大分块处理结果不对频域乘法顺序错检查复数乘法mask失效频域处理mask分段处理5.5 独家避坑技巧技巧一先用PyTorch FFT做原型验证精度和加速比再决定要不要上cuFFT。PyTorch FFT虽然慢但开发快能快速排除不可行方案。技巧二FFT卷积的频域权重可以预计算并缓存推理时直接用。这样省去每次的FFT开销。技巧三如果卷积后面接BN可以把BN的scale和shift吸收到频域乘法里。具体做法是把BN的scale乘到频域权重上shift在IFFT后加。技巧四attention的FFT优化要注意head dimension。如果d不是2的幂padding会引入额外计算。可以尝试把多个head合并成一个大的FFT提高效率。技巧五GEMM的FFT优化要先做结构检测。用SVD或者结构分析工具看一眼权重矩阵确认有Toeplitz或循环结构再用。6. 影响范围与适用边界分析FFT优化其他kernel和算子这件事影响范围比很多人想象的要广但也不是没有边界。从硬件角度看FFT方法在GPU上收益最明显因为GPU的并行能力适合FFT的批处理。在CPU上FFT也有优化但收益通常不如GPU。在专用加速器上要看是否支持FFT指令。从模型角度看FFT方法适合卷积核较大、序列较长、矩阵结构化的场景。对于小卷积核、短序列、稠密矩阵FFT没有优势。从精度角度看FFT方法在fp32下基本无损在fp16下有轻微损失在int8下目前还不成熟。从工程角度看FFT方法的开发成本比标准方法高需要处理padding、内存布局、数值稳定性等问题。所以建议只在收益明显的场景用。我个人的经验是先profiling找到真正的瓶颈再决定要不要用FFT。不要一上来就全面FFT化那样很可能得不偿失。先从一两个关键层试点验证收益后再推广。最后分享一个小技巧如果你在做推理优化可以先把模型里的所有卷积层按卷积核大小排序对大于7×7的层优先尝试FFT。attention层则看序列长度L1024时可以考虑。GEMM层先做结构检测。这样能快速找到FFT的甜点区避免盲目优化。