为何RTX A5000/6000上Torch GPU FFT远慢于CPU？

环境：搭载多块Nvidia RTX A5000/6000显卡、CUDA版本11.8的高性能服务器
现象：运行以下代码后，GPU版本的FFT运算速度比CPU慢200-800倍，更换不同显卡测试结果一致。

测试代码：

import sigpy as sp
import torch
import time

arr = sp.shepp_logan((256, 256))
device = "cpu"
arr = torch.from_numpy(arr).to(device)
tic = time.perf_counter()
res = torch.fft.fft2(arr, dim=(-2, -1))
toc = time.perf_counter()
cpu_time = toc - tic
device = "cuda:5"
arr = arr.to(device)
tic = time.perf_counter()
res = torch.fft.fft2(arr, dim=(-2, -1))
toc = time.perf_counter()
gpu_time = toc - tic
print(f"CPU time: {cpu_time}, GPU time: {gpu_time} ratio: {gpu_time / cpu_time}")

原因分析 #

• CUDA首次初始化开销：第一次调用CUDA设备时，会触发CUDA上下文初始化，包含驱动加载、设备资源分配等操作，这部分开销极大，完全覆盖了FFT本身的计算耗时。而CPU版本无此初始化成本。

• 数据规模与GPU特性不匹配：256×256的矩阵属于小数据量，GPU的优势在于大规模并行计算。小数据场景下，数据从CPU拷贝到GPU的传输时间、GPU核函数启动开销，远大于GPU计算本身的时间；而CPU处理小矩阵FFT的效率天然更高。

• 异步操作导致计时不准确：CUDA操作默认异步执行，原代码中time.perf_counter()计时时，GPU的FFT可能尚未真正完成，且首次初始化的异步过程被错误计入了GPU计算时间，导致计时结果失真。

修正后的测试代码 #

import sigpy as sp
import torch
import time

# 预热CUDA，消除首次初始化开销
torch.cuda.init()
torch.randn(1, device="cuda:5").fft2()
torch.cuda.synchronize()

arr = sp.shepp_logan((256, 256))
arr_cpu = torch.from_numpy(arr).to("cpu")

# CPU测试：多次运行取平均，消除单次波动
cpu_times = []
for _ in range(100):
    tic = time.perf_counter()
    _ = torch.fft.fft2(arr_cpu, dim=(-2, -1))
    toc = time.perf_counter()
    cpu_times.append(toc - tic)
cpu_avg_time = sum(cpu_times) / len(cpu_times)

# 数据传输到GPU，单独分离传输耗时
arr_gpu = arr_cpu.to("cuda:5")
torch.cuda.synchronize()  # 等待传输完成

# GPU测试：多次运行取平均，同步等待确保计算完成
gpu_times = []
for _ in range(100):
    tic = time.perf_counter()
    _ = torch.fft.fft2(arr_gpu, dim=(-2, -1))
    torch.cuda.synchronize()  # 强制等待GPU计算结束
    toc = time.perf_counter()
    gpu_times.append(toc - tic)
gpu_avg_time = sum(gpu_times) / len(gpu_times)

print(f"CPU平均耗时: {cpu_avg_time:.6f}s, GPU平均耗时: {gpu_avg_time:.6f}s, 耗时比: {gpu_avg_time / cpu_avg_time:.2f}")

补充建议 #

• 测试前关闭其他占用GPU资源的进程，保证测试环境纯净。
• 当测试更大尺寸的矩阵（如2048×2048），GPU的并行优势会充分体现，此时运算速度会远超CPU。
• 若需频繁处理小尺寸FFT，建议将多个小矩阵合并为大批次进行计算，最大化利用GPU的并行计算能力。

内容的提问来源于stack exchange，提问作者MRm