fft3d / README.md

README: the 2026-10-27 fixes

dc1c9f3 verified 12 days ago

15.4 kB

---
library_name: kernels
license: MIT
tags:
  - cuda
  - cutlass
  - cute-dsl
  - fft
  - fp32
  - fp64
  - blackwell
  - rtx-5090
---
# 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

```python
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](#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:

| grid (nz x ny x nx) | r2c fp32 | c2r fp32 | r2c fp64 | c2r fp64 |
|---|---:|---:|---:|---:|
| 32x32x32 | 3.33x | 2.60x | 3.93x | 2.22x |
| 64x64x64 | 2.13x | 2.00x | 2.48x | 2.40x |
| 80x80x112 | 2.22x | 1.85x | 2.63x | 2.67x |
| 96x96x96 | 2.24x | 2.14x | 2.43x | 2.52x |
| 96x96x144 | 1.92x | 1.77x | 2.13x | 2.39x |
| 128x128x128 | 1.80x | 1.65x | 2.51x | 2.91x |
| 192x192x192 | 1.39x | 1.44x | 2.75x | 2.95x |
| 256x256x256 | 1.42x | 1.59x | 2.47x | 2.65x |
| geomean | 1.99x | 1.85x | 2.63x | 2.58x |
| geomean, 128^3 and up | 1.53x | 1.56x | 2.58x | 2.83x |

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

| grid | r2c fp32 | c2r fp32 | r2c fp64 | c2r fp64 |
|---|---|---|---|---|
| 32x32x32 | 10.4 / 8.6 / 2.6 | 9.3 / 7.0 / 2.7 | 26.4 / 26.9 / 6.8 | 20.5 / 19.1 / 8.6 |
| 64x64x64 | 11.6 / 9.4 / 4.4 | 14.3 / 9.0 / 4.5 | 35.3 / 37.4 / 15.1 | 41.5 / 37.5 / 15.6 |
| 80x80x112 | 21.1 / 15.1 / 6.8 | 20.1 / 13.9 / 7.5 | 98.7 / 92.7 / 35.2 | 83.6 / 89.7 / 33.6 |
| 96x96x96 | 23.4 / 16.5 / 7.4 | 25.6 / 16.7 / 7.8 | 96.4 / 88.8 / 36.5 | 96.2 / 89.4 / 35.5 |
| 96x96x144 | 31.4 / 18.5 / 9.6 | 31.4 / 18.7 / 10.6 | 128.1 / 115.8 / 54.5 | 146.0 / 116.1 / 48.6 |
| 128x128x128 | 42.0 / 24.2 / 13.4 | 56.8 / 23.4 / 14.2 | 201.4 / 193.9 / 77.1 | 233.7 / 205.0 / 70.4 |
| 192x192x192 | 123.2 / 66.9 / 48.1 | 126.4 / 71.2 / 49.3 | 701.7 / 642.8 / 233.5 | 682.8 / 634.7 / 215.0 |
| 256x256x256 | 346.8 / 163.0 / 114.5 | 300.2 / 187.4 / 117.5 | 1688.8 / 1323.1 / 536.0 | 1671.9 / 1337.3 / 505.1 |

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](#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:

| sample | r2c fp32 | c2r fp32 | r2c fp64 | c2r fp64 |
|---|---:|---:|---:|---:|
| 60 uniform | 1.81x (1.15x) | 1.63x (1.11x) | 1.93x (1.15x) | 1.67x (1.18x) |
| 60 with 7-smooth axes | 1.49x (1.06x) | 1.81x (1.28x) | 1.83x (1.29x) | 1.76x (1.17x) |
| 40 with a prime factor >= 17 | 1.72x (1.22x) | 1.62x (1.15x) | 1.88x (1.20x) | 1.71x (1.16x) |

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`):

| | eight grids (geomean) | 8x8x16 to 48x40x36 (slowest) |
|---|---:|---:|
| r2c fp32 | 1.91x | 1.54x |
| c2r fp32 | 1.72x | 1.11x |
| r2c fp64 | 2.20x | 1.29x |
| c2r fp64 | 2.23x | 1.12x |

## 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](https://github.com/NVlabs/kda/pull/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](#changes)).
MIT-licensed.

## Available functions
- `c2r`
- `prepare_c2r`
- `prepare_r2c`
- `r2c`

	---
	library_name: kernels
	license: MIT
	tags:
	- cuda
	- cutlass
	- cute-dsl
	- fft
	- fp32
	- fp64
	- blackwell
	- rtx-5090
	---

	# 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

	```python
	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](#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:

	\| grid (nz x ny x nx) \| r2c fp32 \| c2r fp32 \| r2c fp64 \| c2r fp64 \|
	\|---\|---:\|---:\|---:\|---:\|
	\| 32x32x32 \| 3.33x \| 2.60x \| 3.93x \| 2.22x \|
	\| 64x64x64 \| 2.13x \| 2.00x \| 2.48x \| 2.40x \|
	\| 80x80x112 \| 2.22x \| 1.85x \| 2.63x \| 2.67x \|
	\| 96x96x96 \| 2.24x \| 2.14x \| 2.43x \| 2.52x \|
	\| 96x96x144 \| 1.92x \| 1.77x \| 2.13x \| 2.39x \|
	\| 128x128x128 \| 1.80x \| 1.65x \| 2.51x \| 2.91x \|
	\| 192x192x192 \| 1.39x \| 1.44x \| 2.75x \| 2.95x \|
	\| 256x256x256 \| 1.42x \| 1.59x \| 2.47x \| 2.65x \|
	\| geomean \| 1.99x \| 1.85x \| 2.63x \| 2.58x \|
	\| geomean, 128^3 and up \| 1.53x \| 1.56x \| 2.58x \| 2.83x \|

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

	\| grid \| r2c fp32 \| c2r fp32 \| r2c fp64 \| c2r fp64 \|
	\|---\|---\|---\|---\|---\|
	\| 32x32x32 \| 10.4 / 8.6 / 2.6 \| 9.3 / 7.0 / 2.7 \| 26.4 / 26.9 / 6.8 \| 20.5 / 19.1 / 8.6 \|
	\| 64x64x64 \| 11.6 / 9.4 / 4.4 \| 14.3 / 9.0 / 4.5 \| 35.3 / 37.4 / 15.1 \| 41.5 / 37.5 / 15.6 \|
	\| 80x80x112 \| 21.1 / 15.1 / 6.8 \| 20.1 / 13.9 / 7.5 \| 98.7 / 92.7 / 35.2 \| 83.6 / 89.7 / 33.6 \|
	\| 96x96x96 \| 23.4 / 16.5 / 7.4 \| 25.6 / 16.7 / 7.8 \| 96.4 / 88.8 / 36.5 \| 96.2 / 89.4 / 35.5 \|
	\| 96x96x144 \| 31.4 / 18.5 / 9.6 \| 31.4 / 18.7 / 10.6 \| 128.1 / 115.8 / 54.5 \| 146.0 / 116.1 / 48.6 \|
	\| 128x128x128 \| 42.0 / 24.2 / 13.4 \| 56.8 / 23.4 / 14.2 \| 201.4 / 193.9 / 77.1 \| 233.7 / 205.0 / 70.4 \|
	\| 192x192x192 \| 123.2 / 66.9 / 48.1 \| 126.4 / 71.2 / 49.3 \| 701.7 / 642.8 / 233.5 \| 682.8 / 634.7 / 215.0 \|
	\| 256x256x256 \| 346.8 / 163.0 / 114.5 \| 300.2 / 187.4 / 117.5 \| 1688.8 / 1323.1 / 536.0 \| 1671.9 / 1337.3 / 505.1 \|

	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](#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:

	\| sample \| r2c fp32 \| c2r fp32 \| r2c fp64 \| c2r fp64 \|
	\|---\|---:\|---:\|---:\|---:\|
	\| 60 uniform \| 1.81x (1.15x) \| 1.63x (1.11x) \| 1.93x (1.15x) \| 1.67x (1.18x) \|
	\| 60 with 7-smooth axes \| 1.49x (1.06x) \| 1.81x (1.28x) \| 1.83x (1.29x) \| 1.76x (1.17x) \|
	\| 40 with a prime factor >= 17 \| 1.72x (1.22x) \| 1.62x (1.15x) \| 1.88x (1.20x) \| 1.71x (1.16x) \|

	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`):

	\| \| eight grids (geomean) \| 8x8x16 to 48x40x36 (slowest) \|
	\|---\|---:\|---:\|
	\| r2c fp32 \| 1.91x \| 1.54x \|
	\| c2r fp32 \| 1.72x \| 1.11x \|
	\| r2c fp64 \| 2.20x \| 1.29x \|
	\| c2r fp64 \| 2.23x \| 1.12x \|

	## 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](https://github.com/NVlabs/kda/pull/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](#changes)).
	MIT-licensed.

	## Available functions
	- `c2r`
	- `prepare_c2r`
	- `prepare_r2c`
	- `r2c`