KDA 3D FFTs on Packed Hermitian Half-Spectra

CuTe-DSL forward (real-to-complex) and inverse (complex-to-real) unnormalized 3D FFTs of one real grid, in fp32 and fp64, for the NVIDIA RTX 5090. The spectrum is the packed half-spectrum that structured-grid spectral solvers use: x innermost, interleaved (re, im), an x-row stride of exactly nx // 2 + 1. The results are those of torch.fft.rfftn and torch.fft.irfftn(..., norm="forward").

Against cuFFT's own 3D plans, geometric mean over eight grids from 32^3 to 256^3: r2c 1.99x and c2r 1.85x in fp32, r2c 2.62x and c2r 2.63x in fp64. Against torch.fft, which runs a 3D transform as two cuFFT plans and a copy, about 3x.

Requirements

• NVIDIA GPU with compute capability 12.0; tuned and tested on the RTX 5090 only
• CUDA 13 compatible driver
• PyTorch 2.12.1 with CUDA 13.0 (tested)
• nvidia-cutlass-dsl[cu13]==4.6.0 (tested)

Usage

import torch
from kernels import get_kernel

fft3d = get_kernel("humanfia-lab/fft3d", version=1, trust_remote_code=True)

u = torch.randn(96, 96, 144, device="cuda")  # [nz, ny, nx], fp32 or fp64
spectrum = fft3d.r2c(u)  # [96, 96, 73, 2] == torch.view_as_real(torch.fft.rfftn(u))
back = fft3d.c2r(spectrum)  # [96, 96, 144], unnormalized: 96 * 96 * 144 * u

launch = fft3d.prepare_r2c(u)  # plan once for this buffer...
u.copy_(torch.randn_like(u))
spectrum = launch()  # ...then transform its current contents

• r2c(x) takes a contiguous [nz, ny, nx] CUDA tensor, float32 or float64. c2r(X) takes a [nz, ny, nf, 2] half-spectrum of the same dtypes and returns [nz, ny, 2 * (nf - 1)]. Each call returns a new tensor and leaves its input unchanged. Other dtypes, devices and shapes raise TypeError or ValueError.
• prepare_r2c(x) and prepare_c2r(X) plan for that input buffer and return launch(), which transforms the buffer's current contents into one output tensor that it reuses on every call. Calls of one launch must be ordered (one stream). Keep launch alive as long as a CUDA graph captured from it. fp64 r2c has no two-phase form; its launch() calls r2c(x).
• Calls on different CUDA streams may run concurrently, including calls of the same shape.
• The first call for a new shape compiles its kernels (CuTe DSL JIT), typically in 1 to 30 s, up to about 4 minutes for a very long z axis over a tiny y-x plane (such as 4096x8x16). Make that first call outside CUDA graph capture; later calls can be captured.
• c2r expects a Hermitian-consistent half-spectrum, such as the output of r2c or Fourier-space data that keeps its symmetry. On other input the result depends on the order of the 1D passes, as with cuFFT; these kernels then match torch.fft.irfftn.
• Tuned on the eight grids below and checked on dozens of others per kernel, including odd and prime extents and up to 1024x1024x512. fp64 r2c raises NotImplementedError when a y or z axis needs a direct DFT longer than 6336 points (prime factors above 13, or smooth lengths above 1536); fp64 c2r raises it for nz from 1280 to 4096, ny from 3072 to 4096 or nx / 2 from 3072 to 4096.

Accuracy

fp32 kernels compute in fp32 from twiddle tables computed in float64 on the host: no reduced precision, no tensor cores, no approximate math. Against the float64 result the relative L2 error is 1.4e-7 to 1.8e-7 on the grids below, where cuFFT's is 2.0e-7 to 2.8e-7. Every element is within 1e-2, or within 1e-3 of its own magnitude.

fp64 kernels are IEEE float64 throughout, within 3.5e-16 to 5.9e-16 relative L2 of torch.fft's float64 result. The RTX 5090 runs FP64 at 1/64 of the FP32 rate, so both fp64 kernels compute part of their float64 additions with a software IEEE add (round to nearest even) on the integer units, which run alongside the FP64 unit. It gives bit-identical results to the FP64 unit for finite normal or zero operands and results; only the sign of some zero results differs. Outside that domain it is wrong, silently: data below about 1e-300 in magnitude, subnormal values, results that overflow, and Inf or NaN inputs give wrong outputs instead of propagating. Data between about 1e-300 and 1e300 in magnitude is unaffected.

Repeated calls give bitwise identical results.

Performance

Speedup over cuFFT's 3D plan:

Times in microseconds (torch.fft / cuFFT 3D plan / kernel):

On small grids cuFFT's several kernels are launch-bound, and one or two fused kernels win most. On large fp32 grids the kernels make fewer passes over memory. In fp64 the FP64 unit is the limit, and the integer-unit adder adds throughput beside it.

Measurements use one RTX 5090 (driver 580.82.07) with CUPTI cold-L2 kernel time over CUDA graph replay, 3 warmup and 30 measured iterations, median of 3 trials, each trial timing cuFFT before and after the kernel. cuFFT is the cuFFT 12.0 that PyTorch 2.12.1+cu130 ships, as one cufftPlan3d per grid made before timing and executed out of place; its C2R overwrites its input and is given a copy. torch.fft is rfftn / irfftn on the same library. Speeds differ by a few percent between RTX 5090 cards.

Source and licensing

Written by Claude Opus 5.5 in four 12-hour runs of KDA (kernel design agents), one per KDA task (fft3d_r2c, fft3d_c2r, fft3d_r2c_f64, fft3d_c2r_f64), built from the requests in NVlabs/kda#7, and accepted by the KDA judge. Published from commit e539498 of humanfia/kda-kernels-internal, directory rtx5090/fft3d/, which also holds its tests and the fixes made after the judge's checks. MIT-licensed.

Available functions

• c2r
• prepare_c2r
• prepare_r2c
• r2c