jburtoft/fnet-neuron-kernels

DEPRECATED -- moved to jburtoft/fnet-fast-fourier-transform-neuron-kernels

This repository is deprecated and frozen at v1. It is kept only so existing KernelConfig references keep resolving.

New location: jburtoft/fnet-fast-fourier-transform-neuron-kernels

Why the move

1. Repo type. This is a model repo, so it never appeared at huggingface.co/kernels and does not resolve through get_kernel. The new repository is a first-class kernel-type repo.
2. Name. "Fast Fourier transform" is now spelled out in the repository name.
3. v2 is substantially faster and more accurate -- see below.

What changed in v2

Area	Improvement
Tensor Engine matmuls	25-50% fewer per tile
Total ISA operations	36-45% fewer per tile
Numerical accuracy	~84x better end-to-end (2.6e-05 -> 3.1e-07 relative)

Principal changes:

• r2c/c2r specialization. FNet's pass 1 has a real input and pass 2 needs only a real output. v1 ran a full complex-to-complex transform in both passes, so half its matmuls either multiplied a zero imaginary input or produced an imaginary result that was discarded.
• No DFT-matrix transpose. W[k,n] = exp(-2*pi*i*k*n/N) is exactly symmetric. Reducing the integer product k*n modulo N before scaling makes the float32 matrix symmetric to 0 ULP, removing 2 nc_transpose plus 2 PSUM copies per DFT.
• Cached constants. v1 rebuilt the DFT and twiddle matrices in NumPy and re-uploaded them on every call -- 144 separate host-to-device transfers for a single B=4, S=512, D=768 forward pass.
• Batch folding instead of a serial Python loop over the batch.
• Real-only final butterfly and fused output assembly.

Migration

The public class was renamed. Update the KernelConfig value:

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

Other renames:

v1	v2
NeuronFNetFourierForward	NeuronFNetFastFourierTransformForward
NeuronFNetFourierLayout	NeuronFNetFastFourierTransformLayout
nki_fft2d_real	nki_fast_fourier_transform_2d_real
_fft1d_{128,256,512}	_fast_fourier_transform_{128,256,512}_{r2c,c2r}

The input and output contract is unchanged: (B, S, D) in, (B, S, D) out, equivalent to torch.fft.fftn(x, dim=(1, 2)).real.

Correction to the v1 documentation

The v1 README stated that non-power-of-two hidden sizes were "padded to the next power of 2 (1024) and the result truncated." That description was wrong. Zero-padding an FFT changes the frequency grid, so truncating a 1024-point transform does not yield a 768-point transform (the two differ by ~116 in absolute terms on unit-scale input).

The v1 implementation was correct -- it computed a true N-point DFT via _torch_fft1d. Only the documentation was wrong. v2 documents the actual behavior and asserts it in test_e2e.py.

Relatedly, v1 reported accuracy as "cosine similarity 1.000001-1.000007". Cosine similarity cannot exceed 1.0, and it is insensitive to the broadband errors this kernel can produce. v2 reports max_abs_err against a float64 torch.fft reference instead.

License

Apache 2.0