torch_cufft: Extending PyTorch FFT Capacity with Multi-GPU and Descriptor-Resident cuFFTXt Execution
TL;DR
This extended abstract presents torch_cufft, a PyTorch-facing FFT system that routes FFT requests across single-GPU, single-node multi-GPU, and distributed multi-process backends, and focuses on a simple systems question: how much does the maximum feasible FFT size depend on where data lives during execution.
Abstract
Large scientific images and spectral-learning workloads often require two-dimensional Fast Fourier Transforms (2D FFTs) that exceed single-GPU memory. This matters for Fourier Neural Operators (FNOs) and Transform Once (T1)-style models, where frequency-domain computation is central to the learning workflow. PyTorch provides convenient FFT APIs, but scaling these transforms across multiple GPUs requires lower-level libraries and careful memory-layout management. This extended abstract presents torch_cufft, a PyTorch-facing FFT system that routes FFT requests across single-GPU, single-node multi-GPU, and distributed multi-process backends. We focus on a simple systems question: how much does the maximum feasible FFT size depend on where data lives during execution? In single-node capacity experiments, descriptor-resident cuFFTXt execution extends the 2D FFT frontier beyond both single-GPU PyTorch and integrated sharded XT execution, reaching 74752 × 74752 on 4 × V100 and 88064 × 88064 on 4 × L40S while preserving FP32-scale correctness in validated configurations.