KDA Batched 1D Complex FFT

CuTe-DSL forward FFT of each row of a batch of complex fp32 rows, the same result as torch.fft.fft, for the NVIDIA RTX 5090:

y[b, f] = sum over t of x[b, t] * exp(-2j * pi * f * t / n)

It handles any length: smooth lengths with in-shared-memory and L2-resident multi-pass kernels, primes with Rader's algorithm and everything else with Bluestein's.

1.84x geometric-mean speedup over cuFFT (torch.fft.fft) and 1.60x over VkFFT across 17 workloads; about 2x on large power-of-two rows, where the kernel runs at the speed of a device-to-device copy.

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

fft = get_kernel("humanfia-lab/fft-c2c-1d", version=1, trust_remote_code=True)

x = torch.randn(32, 1 << 22, 2, device="cuda")  # [batch, n, (re, im)], fp32
y = fft.run(x)  # torch.view_as_real(torch.fft.fft(torch.view_as_complex(x)))

• Any n >= 1 and any batch. The transform is forward and unnormalized. The result is a new [batch, n, 2] fp32 tensor; x is left unchanged.
• The first call for a new (batch, n) compiles its kernels (CuTe DSL JIT); later calls reuse them. Make that first call outside CUDA graph capture; later calls can be captured.
• Calls on different CUDA streams may run concurrently.

Accuracy

Everything is computed in fp32: no reduced precision, no tensor cores, no approximate math. Twiddle factors come from tables computed in float64 on the host.

On every workload below, every output element is within 1e-5 of the float64 result, relative to the tensor's RMS or to the element itself. The relative L2 error is 1.1e-7 to 3.4e-7. On the same workloads it is 1.4e-7 to 4.1e-7 for cuFFT, 3.0e-7 to 1.1e-6 for VkFFT and 1.1e-7 to 3.2e-7 for VkFFT with its lookup tables. Repeated calls give bitwise identical results.

Performance

On large rows of smooth length the kernel reads and writes every element once: at n = 2^20, batch 128, it moves 1.56 TB/s, as fast as a device-to-device copy of the same tensor. cuFFT and VkFFT pass over DRAM twice there. For n = 16384, 400, 1000 and 3072, cuFFT already passes once, so the kernel only matches 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 the baseline before and after the kernel. cuFFT is torch.fft.fft with the cuFFT 12.0 that PyTorch 2.12.1+cu130 ships; the newest cuFFT, 12.4.0.43, takes the same time to within 0.5%, except for n = 4093, where it is slower. VkFFT is pyvkfft 2025.1.1 (CUDA backend, out of place), the faster of its default and lookup-table configurations. Speeds differ by a few percent between RTX 5090 cards.

Source and licensing

Written by Claude Opus 5.5 in a 12-hour run of KDA (kernel design agents) for the task fft_c2c_1d, and accepted by the KDA judge. Published from commit 11b486e of humanfia/kda-kernels-internal, directory rtx5090/fft_c2c_1d/, which also holds its tests. MIT-licensed.

Available functions

• run