stft-mel

Fused STFT to mel spectrogram in a single pass, loadable through kernels. Framing, window, real FFT, power spectrum, mel projection, and log collapse into one traversal, so no complex spectrogram is ever materialized. The reference baseline is torch.stft plus a dense einsum mel projection, and a float64 STFT and projection for accuracy.

The Whisper front end runs before every transcription, and on a CPU it is one of the few remaining places where a fixed, well-understood transform still costs milliseconds per utterance. This kernel computes it as a single vertical pass across frames: no lane-crossing shuffle anywhere in the FFT, no complex intermediate written to memory, and a banded mel projection that touches only the nonzeros. Thirty seconds of audio becomes 80 mel bands in about two milliseconds.

The kernel's own output: 30 s of a synthetic voiced signal at 16 kHz transformed to 80 log-mel bands, painted as the transform advances, in 2.1 ms end to end.

Usage

import kernels

sm = kernels.get_kernel("phanerozoic/stft-mel", version=1, trust_remote_code=True)

mel = sm.MelSpectrogram(sr=16000, n_fft=400, hop=160, n_mels=80)
feats = mel(audio)          # [B, 80, n_frames] log-mel, Whisper's front end

MelSpectrogram caches the window and filterbank; reuse one instance across calls. log_mel(audio, ...) is the one-shot form, and spectrogram(...) returns magnitude or power without the mel stage. version selects the release branch; trust_remote_code is required by kernels for publishers without the trusted-publisher mark.

API

Method

The transform is vectorized across frames, not within one: lane v of every SIMD register holds frame f0+v, so each butterfly is a vertical operation and no lane-crossing shuffle occurs anywhere in the FFT. Twiddles are identical across lanes and broadcast from a table. STFT always has far more frames than lanes, which makes this the natural axis and gives the scalar dataflow widened by the vector length with no restructuring.

The real transform of size n_fft runs as a complex Stockham autosort FFT of size n_fft/2 plus a split step, over mixed radices 2/3/4/5. Each pass is instantiated at a compile-time radix, which removes the branch from the inner loop and makes the twiddle index exact without a modulo. Mel filters are triangular, so each output bin reads a short contiguous run of frequency bins; the bank is stored banded as (start, length, weights) and the projection costs O(nnz) rather than O(n_mels * n_freq), under a quarter of the dense weights at n_fft=400, n_mels=80. On x86-64 the binary carries both an AVX2 and a baseline lowering behind an IFUNC resolver, selected at load time; compiling for a single ISA measured 1.85x slower on non-native hosts.

Measured

30 s of mono audio, best of 20 runs, against torch.stft followed by a dense einsum mel projection; both produce log-mel:

The Whisper front end runs 30 s of audio in 2.32 ms on x86-64 and 7.02 ms on a Raspberry Pi 5.

Accuracy

Single precision throughout. Against a float64 STFT and float64 mel projection, the maximum relative error over a well-conditioned bank is within a small factor of what torch.stft in float32 achieves on the same input, and the test suite asserts that bound rather than comparing two float32 implementations to each other. Banks with more mel bins than the FFT can resolve, such as n_mels=128 at n_fft=400, produce filters one or two bins wide; a mel value is then a single float32 product and relative error reflects float32 conditioning rather than the transform. Both this kernel and torch.stft show it.

Requirements and limits

• n_fft must be even, and n_fft/2 must factor over {2, 3, 5}. That covers every value used in audio: 400, 480, 512, 800, 960, 1024, 2048. Anything else raises rather than falling back silently.
• CPU, x86_64-linux and aarch64-linux, for PyTorch 2.11, 2.12, and 2.13 (six prebuilt variants). Each architecture is built on its own host and validated against the prebuilt binary.

References

Stockham autosort FFT; the Whisper log-mel front end as the target configuration; triangular mel filterbank construction.

License

Apache-2.0.