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.63x and c2r 2.58x in fp64. Against torch.fft, which runs a 3D transform as two cuFFT plans and a copy, about 3x. Any grid runs FFT kernels: on random grids up to 256 per axis, 1.5x to 1.9x over cuFFT (geomean), and no sampled grid slower than cuFFT.

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. The work runs on the input's device (which need not be the current one), on its current CUDA stream.
• 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). A launch() on another stream than the one prepare_* ran on waits for prepare_*'s work and keeps its output and workspaces reserved for that stream's work, even if the launch object is dropped before the work ran.
• Calls on different CUDA streams may run concurrently, including calls of the same shape. Calls may also come from several threads: CuTe DSL is not thread-safe while it compiles or makes the first call of a compiled function, so a new shape's plan is built, and makes its first call, under a lock that all kda kernels share (later calls take no lock).
• A call costs about 13 to 17 us of host time (TVM-FFI launches on the current stream). Timed per call as the requests' benchmark does (flashinfer-bench: synchronize, event, call, event), every kernel is faster than torch.fft on every grid tried, the smallest included (see Performance).
• The first call for a new shape compiles its kernels (CuTe DSL JIT), typically in 1 to 15 s, up to about 4 minutes for a very long z axis over a tiny y-x plane (such as 4096x8x16). That first call may itself be inside CUDA graph capture.
• 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; any grid with nz, ny >= 8 and nx >= 16 is accepted. Within 256 per axis every grid of every kernel runs FFT kernels, odd and prime extents included, checked on RTX 5090s by sweeps over every plan signature, and on grids up to 1024x1024x512. Beyond 256 per axis, axis lengths with two prime factors above 16 (fp32 c2r) or a squared prime factor above 13 (fp64, such as 289) still take slow generic kernels. fp64 r2c raises NotImplementedError when a y or z axis needs a direct DFT longer than 6336 points; fp64 c2r raises it for nz from 1280 to 4096, ny from 3072 to 4096 or nx / 2 from 3072 to 4096. The fp64 kernels address with 32-bit offsets: fp64 r2c raises ValueError past 2^30 complex outputs and fp64 c2r on 2^31 or more half-spectrum elements.

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. On other grids within 256 per axis the relative L2 error is 1e-7 to 3.4e-7, every element passes that rule on random input, and large prime extents stay below cuFFT's error (c2r of a random 251x251x254 half-spectrum: max 0.006, cuFFT 0.016).

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 returns exactly what the FP64 unit returns for every input, subnormal operands and results, overflow, signed zeros, Inf and NaN included (2.5e10 operand pairs per routine compared with the native instructions). Until 2026-10-24 it was exact only for finite normal operands and results, and data near the bottom of the normal range came out wrong (see Changes). Data with a large mean keeps cuFFT-level accuracy or better in every kernel.

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. The fp64 columns were measured again after the IEEE-complete adders (see Changes), on an RTX 5090 that runs both cuFFT and the kernels 10-20% slower in absolute time; the speedups moved by at most 2%. Speeds differ by a few percent between RTX 5090 cards.

Beyond the eight grids, against cuFFT's 3D plans on fixed seeded samples of the requests' range (nz, ny 8..256, nx 16..256, even nx for c2r; random input), geomean with the slowest grid in brackets:

Every sampled grid is faster than cuFFT's 3D plan and than torch.fft. Large prime extents gain most: 251x251x254 takes 0.25 ms in fp32 r2c (cuFFT 0.59 ms) and 2.6 ms in fp64 c2r (cuFFT 4.9 ms).

Timed per call the way the requests' benchmark (flashinfer-bench 0.1.2) does it, on fresh copies of the input with a cleared L2 and the host's launch work included, against the requests' baseline (torch.fft):

Changes

Updated on 2026-10-24 with fixes found by sweeping every planner signature over nz, ny in 8..256 and nx in 16..256 (the requests' range), on a B300 with the RTX 5090's plans:

• fp32 r2c returned garbage, often NaN, on 11,765 grids of that range, all with nx >= 84 and ny >= 84 (such as 8x126x140, 8x252x156, 70x104x224): their plane or y-split plan needs a whole plane or sub-plane as one row block, and the x pass halved the rows that did not fit in shared memory. Those grids now take another plan; the plans of the eight grids in the tables above are unchanged.
• The fp64 software adders were exact only for finite normal operands and results: fp64 r2c of a 96x96x144 grid with a quarter of its z planes at 1e-308 to 1e-300 returned NaN everywhere, and data scaled by 1e-300 or 1e-310, subnormal bands of a spectrum and single NaN or Inf inputs failed on many grids. They now handle every IEEE case, and fp64 c2r no longer routes 2x - y through software on five tuned grids. On an RTX 5090 the eight grids' speedups moved by at most 2% (the fp64 columns above).
• fp64 inputs on a CUDA device other than the current one raised (ValueError, BufferError); every call now runs under the input's device.

Updated again on 2026-10-24 after a sweep on RTX 5090s over every plan signature of the four kernels within 256 per axis (about 45,000 grids, random and Hermitian c2r input, spectrum-shaped and offset data, prepare, graph and two-stream checks):

• fp32 c2r missed the 1e-2 / 1e-3 tolerance on about 6,900 grids of 6.5M points and more with a large prime extent (such as 251x251x254) on random half-spectra, the requests' own input: up to 69 elements per grid off by up to 0.03, where cuFFT's fp32 plan has none. Now at most 0.006.
• fp32 r2c returned garbage on grids of 2^24 points and more with an odd ny (such as 256x297x256, beyond 256 per axis).
• Most grids ran slow fallbacks: fp32 r2c ran every grid whose nx is not a multiple of 4 or that has a prime factor above 13 (96.8% of the range) as O(n^2) direct DFTs, 17 to 180 times slower than torch.fft (96x96x98: 0.97 ms against cuFFT's 0.020 ms); fp32 c2r ran 77% of the range on a global-memory generic path, and fp64 ran axes with a prime factor above 13 as direct DFTs (fp64 c2r: the whole grid, about 90% of the range), 3 to 10 times slower than cuFFT. They now run general FFT kernels (two real rows per complex FFT for odd x, Rader's or Bluestein's algorithm for large primes), faster than cuFFT on every sampled grid (Performance). Sweeps over every new kernel signature found no wrong result.
• fp64 r2c on data with a large mean (randn + 1e3) missed atol 1e-6 on 21 of 100 sample grids of 1M points and more, where cuFFT passed; now at most 1.07x cuFFT's error.
• Per-call host time: the fp64 kernels spent about 40 to 50 us of host time per call (0.48x to 0.69x of torch.fft on small grids when timed per call, as the requests' benchmark does). All kernels now launch through TVM-FFI, at about 13 to 17 us; outputs are bitwise unchanged. fp64 c2r now transforms an input that is not 16-byte aligned from an aligned copy instead of raising.

The eight grids keep their plans and their outputs are bitwise unchanged.

Updated on 2026-10-27 with the fixes of a second audit on RTX 5090s and B300s (threads, streams, CUDA graph capture, 32-bit limits; no wrong result was found within 256 per axis):

• Threads making first calls at once failed: CuTe DSL is not thread-safe while it compiles or makes the first call of a TVM-FFI function. Concurrent plan builds raised MLIR errors or ended the process, and two first calls of one shape at once left a spin lock held, so that the next first call of any kernel in the process hung for good. Plans are now built, and make their first call, under one process-wide lock; later calls take none.
• The first call of a shape inside CUDA graph capture raised; its tables are now copied on a stream of the package's own, outside the graph.
• A launch() right after prepare_* on another stream could run before work still queued on prepare_*'s stream, which then overwrote its output, and a launch object dropped before its kernels ran handed its buffers to new tensors on prepare_*'s stream. Such a launch now waits for prepare_*'s stream and keeps its buffers until its work is done; same-stream launches are unchanged.
• fp32 r2c transforms a contiguous input that is not 16-byte aligned from an aligned copy instead of raising.
• fp64 r2c past 2^30 complex outputs (an illegal memory access) and fp64 c2r on 2^31 or more half-spectrum elements (a wrong result) now raise ValueError, beyond 256 per axis.
• fp64 axes of 1529 points and more whose big-prime tile does not fit in shared memory failed to compile (such as 8x6329x16); they keep the direct DFT.

The eight grids keep their plans; their outputs and speeds are unchanged.

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/ (now sm120/fft3d/), which also holds its tests and the fixes made after the judge's checks; updated from commits b0811f8 and 2238a9e (see Changes). MIT-licensed.

Available functions

• c2r
• prepare_c2r
• prepare_r2c
• r2c