FNet Fast Fourier Transform NKI Kernels for AWS Neuron

NeuronCore kernels that compute the 2D fast Fourier transform used by FNet entirely on-device, written in NKI. The 128-point DFT base case runs as a matrix multiply on the Tensor Engine.

Drop-in replacement for FNetBasicFourierTransform.

Why this kernel exists

torch.fft.fftn has no on-device implementation on Neuron -- complex dtypes are not supported by the compiler. Verified still true on PyTorch Native Beta 5 (neuronx-cc 2.27, NKI 0.6); see "Compiler support" below.

Measured on the real 12-layer google/fnet-base, swapping only FNetBasicFourierTransform.forward:

The kernel's value is that it keeps the layer on device, which is what makes the model compilable. With torch.compile it is 8.5x faster than the best eager fallback; on Beta 5 defaults the stock model does not run at all. In plain eager mode with permissive fallback the two are within ~10% -- the transform is not the eager bottleneck at this size.

Usage

from transformers import AutoModel, KernelConfig

kernel_config = KernelConfig({
    "FNetBasicFourierTransform":
        "jburtoft/fnet-fast-fourier-transform-neuron-kernels:NeuronFNetFastFourierTransformForward",
})

model = AutoModel.from_pretrained(
    "google/fnet-base",
    kernel_config=kernel_config,
    device_map="neuron",
)

Or load the kernel directly:

from kernels import get_kernel

k = get_kernel("jburtoft/fnet-fast-fourier-transform-neuron-kernels",
               version=2, trust_remote_code=True)
out = k.nki_fast_fourier_transform_2d_real(hidden_states)   # (B, S, D) -> (B, S, D)

How it works

FNet applies torch.fft.fftn(hidden_states, dim=(1, 2)).real to (B, S, D). This factors into two passes of 1D transforms, each of which is half a complex transform:

Specializing each pass to its actual role halves the Tensor Engine matmul work versus running a full complex-to-complex transform twice.

Each 1D transform uses flat radix-2 Cooley-Tukey:

• 128-point: direct DFT via nc_matmul on the 128x128 Tensor Engine
• 256-point: 2 groups of 128 + 1 butterfly level
• 512-point: 4 groups of 128 + 2 butterfly levels

Performance

End-to-end: real google/fnet-base, 12 layers, on device

This is the number that matters. Measured on the actual transformers FNetModel (12 layers, hidden 768, batch 1, seq 128) with FNetBasicFourierTransform.forward swapped per variant. trn2.3xlarge, PyTorch Native Beta 5 (torch 2.12.1, torch-neuronx 2.12.3.0.1636, neuronx-cc 2.27.2878, NKI 0.6.0), driver 2.30.2, p50.

Three conclusions, in order of importance:

1. Under torch.compile(backend="neuron") the kernel is the only option that works, and it is 8.5x faster than the best eager fallback (3.07 ms vs 26.18 ms). Every torch.fft variant fails to compile with failed to legalize operation 'torch.operator'. The kernel compiles into one graph with zero fallback ops, and compiling is worth 8.9x on the kernel path alone (27.33 -> 3.07 ms). This is where the kernel's value actually lives.
2. In Beta 5's default strict mode, the stock model cannot run at all. torch.fft raises CPU fallback disabled ... Operation aten::_fft_c2c received a complex-dtype input. The kernel is not an optimization here; it is a prerequisite.
3. In eager permissive mode the four variants are within ~10% of each other (23.8-29.0 ms). At FNet-base's size the transform is not the bottleneck in eager mode -- embeddings, 12 dense FFN blocks and LayerNorms dominate. Eager microbenchmarks of the transform alone are therefore a poor guide to model-level impact.

⚠ Correction to the previously published CPU-fallback speedup

Earlier versions of this card claimed 233-1499x faster than the CPU fallback. That claim was wrong and has been withdrawn.

The baseline was torch.fft.fftn(h, dim=(1,2)).real.to(device). .real on a complex tensor returns a non-contiguous view (stride ..., 2), and transferring that view host-to-device degenerates into a pathological element-wise DMA:

Host-to-device transfer of a 393 KB tensor	Time
contiguous	0.139 ms
non-contiguous .real view	686.8 ms
.real.contiguous()	0.178 ms

One .contiguous() call is worth 1685x. The old baseline measured a bad transfer, not the cost of computing an FFT on the host.

The v1-vs-v2 comparison is unaffected (both measured identically, on-device), so the 4.0-9.6x v2-over-v1 improvement still stands.

Isolated transform latency (context, not the headline)

Single transform, no surrounding model. Included for completeness -- and as a caution that these numbers do not predict model-level behavior:

In eager mode at small sizes a host round trip is genuinely faster than this kernel. The kernel is dominated by eager per-tile dispatch (a B1 S128 D768 call issues 7 separate @nki.jit launches), not arithmetic. That penalty disappears under torch.compile, which is why the isolated ranking inverts at the model level.

A queue-pressure sweep (0-16 matmuls queued before the transform, forcing the fallback to drain them) narrowed the gap from 1.66x to 1.20x but did not flip it in eager mode. Pipeline-drain cost is real but modest at this scale; the compile barrier dominates.

When to use this kernel

Use it when:

• You compile with torch.compile(backend="neuron") -- torch.fft cannot be compiled at all, so this is the only way to keep the layer in a compiled graph. Worth 8.5x end-to-end.
• You run on Beta 5 defaults -- the stock path raises.
• You need zero host syncs -- to keep an async pipeline full, or inside a larger fused region.
• Transforms are large (S x D >= 512 x 768), where the kernel wins even in eager mode.

Reach for the host path instead when you are in eager mode, at small transform sizes, with permissive fallback available, and you do not care about host syncs.

v2 vs v1

Unaffected by the baseline correction -- both versions measured identically and entirely on-device. Beta 4, p50 of 30 iters:

v2 is 4.0-9.6x faster than v1.

Where the v2 speedup comes from

The static instruction-count reduction is only ~1.6x at best, so it cannot by itself explain 4-9.6x. Ablation isolates two effects that multiply. Both runs cache constants, so the only difference in the "kernel" column is the r2c/c2r specialization plus the removed W transpose:

• Constant caching: 2.2-4.8x. The largest win, and not a kernel change at all. v1 rebuilt its DFT and twiddle matrices in NumPy and re-uploaded them every call. Instrumented at B=1 S=512 D=768: v1 performs 36 host-to-device transfers totalling 3.15 MB per forward pass; v2 performs 0.
• Kernel specialization: 1.6-2.9x. Larger than the raw matmul count suggests, because removing the W transposes also removes 2 PSUM->SBUF round-trips per DFT, and the c2r path drops the final butterfly's imaginary half entirely.

Do not attribute the full 4-9.6x to the matmul specialization.

Batch folding measured as neutral (within noise) once constants are cached.

Compiler support for torch.fft (Beta 5, neuronx-cc 2.27)

Tested directly on Beta 5. Complex dtypes and the 2D FFT are still not supported on device:

Every torch.fft path that "works" executes on the host, confirmed by torch_neuronx.get_fallback_ops(). The kernel reports no fallback ops.

Beta 5 also sets TORCH_NEURONX_FALLBACK_ONLY_FOR_UNIMPLEMENTED_OPS=1 during import torch, so unsupported ops raise by default:

RuntimeError: CPU fallback disabled because
TORCH_NEURONX_FALLBACK_ONLY_FOR_UNIMPLEMENTED_OPS=1.
Operation aten::_fft_c2c received a complex-dtype input; complex dtypes are unsupported.

Permissive fallback must be re-enabled explicitly (set the variable to 0 after importing torch) for the host path to work at all.

Note: torch.compile requires the kernel to be importable as a real module. A module loaded via importlib.spec_from_file_location without registering it in sys.modules fails under Dynamo with ModuleNotFoundError.

Instruction counts

ISA operations per 128-point tile, counting both FNet passes:

Sources of the reduction:

1. r2c / c2r specialization -- skips matmuls against a zero imaginary input (pass 1) and matmuls producing a discarded imaginary output (pass 2).
2. No DFT-matrix transpose. W[k,n] = exp(-2*pi*i*k*n/N) depends on k*n and is therefore exactly symmetric. Angles are built by reducing the integer product k*n modulo N before scaling, which makes the float32 matrix symmetric to 0 ULP, so W is fed to nc_matmul directly. This removes 2 nc_transpose plus 2 PSUM->SBUF copies per DFT.
3. Real-only final butterfly (c2r path) -- 5 vector ops instead of 10.
4. Fused output assembly -- butterflies write straight into their output halves instead of writing temporaries and copying.
5. Cached constants -- DFT matrices, twiddle factors and -W_imag are built once per (size, device) instead of being rebuilt in NumPy and re-uploaded on every call. Previously a B=4, S=512, D=768 forward pass issued 144 separate host-to-device transfers.
6. Batch folding -- the batch is folded into the tile dimension ((B,S,D) -> (B*S,D)) instead of being looped over in Python.

Supported sizes

The fallback is a genuine N-point DFT, not a zero-padded power-of-two transform. Zero-padding would change the frequency grid and produce a different result.

Note that google/fnet-base has hidden_size=768, which is not a power of two, so its pass 1 runs the dense on-device DFT and only pass 2 uses the NKI kernels. Both paths stay on-device; the transform is never sent to the host.

Accuracy

Validated against torch.fft.fftn(...).real (float64 reference) on trn2.3xlarge hardware (PyTorch Native Beta 4, SDK 2.31). test_e2e.py runs 30 checks; all pass.

v2 is 171-1091x more accurate than v1, measured on the same hardware in the same run. The modular angle reduction is the cause; it costs nothing at runtime.

The gap is widest at D=768, where v1's error reached 2.0e-04 relative. That is the FNet-base hidden size, i.e. the configuration most users will hit.

Note that bfloat16 input yields ~2.9e-03 relative error in both versions. That is inherent to bf16's 8-bit mantissa, not a kernel defect -- the FFT arithmetic is always performed in float32 and only the returned tensor is cast back.

Requirements

• AWS Neuron SDK 2.29+ (NKI 0.3.0+); validated on NKI 0.5.0 / SDK 2.31
• PyTorch Native (torch-neuronx) -- the orchestration layer runs eagerly
• transformers with KernelConfig support, and the kernels library

Framework note: the @nki.jit kernels use only ISA-level nisa.* operations and work under both PyTorch Native and the standard torch-neuronx XLA trace path. The orchestration layer (nki_fast_fourier_transform_2d_real) uses eager tensor operations and therefore needs PyTorch Native. Using these kernels inside an XLA trace requires a trace-compatible orchestration layer.

API

NeuronFNetFastFourierTransformForward     HF forward class (use this in KernelConfig)
NeuronFNetFastFourierTransformLayout      weight layout (empty -- no learnable params)
nki_fast_fourier_transform_2d_real        (B,S,D) -> (B,S,D) real part of 2D FFT

_fast_fourier_transform_{128,256,512}_r2c real -> complex  @nki.jit kernels
_fast_fourier_transform_{128,256,512}_c2r complex -> real  @nki.jit kernels
_dft_r2c / _dft_c2r / _dft_c2c            128-pt DFT primitives via nc_matmul

Changelog

v2.3 -- Documentation-only. Added end-to-end measurements on the real 12-layer google/fnet-base, which is the benchmark that actually answers "is this kernel worth using." Result: under torch.compile(backend="neuron") the kernel runs at 3.07 ms and is the only variant that compiles at all (every torch.fft path fails to legalize), making it 8.5x faster than the best eager fallback. On Beta 5 defaults the stock model raises. In eager permissive mode all four variants are within ~10%. Prior versions of this card only ever measured the transform in isolation, which understated the kernel's value (it looked slower) by omitting the compile barrier and the strict-mode failure.

v2.2 -- Documentation-only correction. Withdrew the "233-1499x faster than the CPU fallback" claim, which was measured against a pathological non-contiguous host-to-device transfer (.real on a complex tensor yields a stride-2 view; sending it H2D costs 686.8 ms vs 0.178 ms with .contiguous(), a 1685x difference). Added Beta 5 compiler-support results confirming complex dtypes and 2D FFT are still unsupported on device. No code changes.

v2.1 -- Hardware benchmark on trn2.3xlarge (PyTorch Native Beta 4, SDK 2.31): 4.0-9.6x faster than v1, 171-1091x more accurate, test_e2e.py 30/30 PASS. Attribution: constant caching 2.2-4.8x (v1 issued 36 H2D transfers per forward, v2 issues 0), kernel specialization 1.6-2.9x, batch folding neutral. Added benchmarks/.

v2 -- r2c/c2r specialization, symmetric-W transpose elimination, real-only final butterfly, fused output assembly, constant caching, batch folding, exact modular angle reduction. Static counts: 36-45% fewer ISA ops, 25-50% fewer matmuls. Public class renamed to NeuronFNetFastFourierTransformForward. Corrected documentation: the previous release's README incorrectly stated that non-power-of-two sizes were zero-padded to the next power of two and truncated; the implementation has always used a true N-point DFT.

v1 -- initial release: complex-to-complex 128/256/512 NKI kernels, NeuronFNetFourierForward.

License

Apache 2.0