flashmel

Fused CUDA mel spectrogram: framing + padding + Hann window + R2C FFT (cuFFTDx)

• |X|^power + sparse mel filterbank projection in a single kernel launch, with a drop-in torchaudio.transforms.MelSpectrogram-compatible Python wrapper.

3–14× faster than torchaudio on an RTX 4070 Laptop, bit-matching torchaudio at atol=rtol=1e-4 (fp32).

Installation

flashmel needs Python ≥ 3.10, an NVIDIA GPU, a recent NVIDIA driver, and a CUDA build of torch (and torchaudio) — install those first, the way your environment needs them, e.g.:

pip install torch torchaudio --index-url https://download.pytorch.org/whl/cu130

flashmel deliberately does not depend on torch/torchaudio: the CUDA wheel stack is environment-specific (driver version, CUDA version, sometimes a private mirror), and pinning it here would only fight your install. Then:

pip install flashmel

This pulls flashmel's own deps — the cuFFTDx/CUDA headers and NVRTC (nvidia-mathdx, nvidia-cuda-cccl, nvidia-cuda-runtime, nvidia-cuda-nvrtc), the driver bindings (cuda-python), and numpy. No system CUDA toolkit or nvcc is required. The kernel is compiled on first use per configuration (~2 s) and cached on disk (see below).

Usage

import torch
from flashmel import MelSpectrogram

mel = MelSpectrogram(n_fft=400, hop_length=160, n_mels=128)  # torchaudio args
x = torch.randn(8, 480000, device="cuda")                    # f32 or i16 CUDA tensor
y = mel(x)                                                   # [8, 128, 3001]

Supports every torchaudio.transforms.MelSpectrogram argument (sample_rate, n_fft, win_length, hop_length, f_min, f_max, pad, n_mels, power, normalized, center, pad_mode, onesided, norm, mel_scale) except window_fn/wkwargs (always periodic Hann). Additionally:

• input dtype: float32, or int16 raw PCM (scaled by 1/32768) — dispatched from the input tensor's dtype;
• output dtype: dtype=torch.float32 (default) or torch.float16;
• input shape (..., time), output (..., n_mels, n_frames), like torchaudio.

Constraints: n_fft must be 400 or a power of 2 in [64, 16384]; n_mels ≤ 256; onesided=True only (torchaudio's own MelScale can't consume two-sided spectrograms either); input must be on CUDA.

How it works

• NVRTC-compiled flashmel/kernel.cu, one cubin per parameter configuration, disk-cached by content hash in the user cache dir (~/.cache/flashmel, overridable via FLASHMEL_CACHE_DIR). No build step; first call per config compiles in ~2 s.
• Frames are loaded straight into cuFFTDx registers with the signal padding (pad, center, all four pad_modes) applied as index arithmetic — the padded/framed signal is never materialized.
• The STFT normalized modes are folded into the window vector on the host.
• Power-of-2 sizes use cuFFTDx real_mode::folded R2C (an N/2-point complex FFT), halving shared memory and butterflies; 400 uses real_mode::normal.
• The FFT's shared-memory workspace is reused as the per-frame power tile (odd row stride ⇒ bank-conflict-free), from which the sparse mel projection reads. Mel filters are triangular ⇒ contiguous support, packed as start[m] / width[m] / weights[m, W_MAX]; the inner loop runs to each filter's actual width, which cut L1/SMEM traffic ~40% for wide-filter configs (1024: 0.28→0.23 ms, 4096: 0.57→0.41 ms, 16384: 1.35→0.94 ms).
• When the block has more threads than mel outputs (large n_fft ⇒ FPB=1), SPLIT lanes cooperate per dot product with a warp-shuffle reduce instead of idling (16384: another 12%).
• Folded-mode interior frames (no padding in play) load signal and window as 2-wide vectors, halving input-phase load instructions (64: −30%, 1024/4096: −13%; the LSU pipe, not DRAM, is the binding resource at these sizes).
• Normal-mode (400) interior frames take a direct-load path that skips the per-sample reflect/boundary arithmetic entirely (whisper case: 1.08→0.54 ms, 2.0×, now at copy peak).
• Launches via the CUDA driver API on torch's current stream; fp16 output is cast at the final store (all accumulation in fp32).
• The launch is wrapped in a torch.library custom op (flashmel::mel_forward with a fake impl), so torch.compile(..., fullgraph=True) traces through it.
• Kernel modules are cached per (config, arch, device) — CUmodules are context-bound — and filter/window buffers are cached per device, so multiple GPUs work concurrently. Inference-only: requires-grad input under grad mode raises (no backward is implemented).

Performance (RTX 4070 Laptop, sm_89, measured copy peak ≈ 198 GB/s)

uv run python scripts/bench.py (CUDA-event timing, median of 50; "BW" counts essential bytes = input + output only):

Full-spectrum sweep (hop = n_fft/4, 128 mel, B=32, T=160000): 64–512 run at 97–120% of copy peak, 1024 at 82%, then declining as the FFT itself dominates. Numbers drift ±15% with laptop thermals.

Where each case actually sits (Nsight Compute)

• small–512: DRAM-saturated. Effective BW exceeds copy peak because overlapped frame reads hit L2 (essential-byte metric counts them once).
• speech (1024) / music (4096): L1/SMEM-throughput bound at 76–81% SOL with DRAM at 30–40% — the remaining L1 traffic is cuFFTDx's internal register exchanges, not under our control (vectorizing our own smem reads is blocked by the odd power-tile stride that keeps them bank-conflict-free).
• whisper (400): DRAM-saturated since the interior direct-load path (DRAM 77% / L1 80% / SM 74% SOL — 77% of theoretical DRAM ≈ the measured copy peak; previously latency-bound at DRAM 48% / SM 47% with the reflect arithmetic on every sample). Occupancy stays at 41%, capped by the FFT workspace (12.8 KB per 4-frame block ⇒ 7 blocks/SM; normal-mode R2C is the only cuFFTDx option for non-power-of-2), but that no longer binds. A paired c2c trick (two frames per complex FFT via conjugate symmetry) halves FFT work but measured ~1.5× slower: the full Z-spectrum round trip through shared memory costs more than the saved butterflies. Implemented, A/B'd, removed.
• huge (8192/16384): occupancy-bound (33–64 KB workspace ⇒ 1–2 blocks/SM) and barrier-serialized inside the block FFT. Inherent to a single-kernel block FFT at these sizes; going faster would mean a multi-kernel FFT (and materializing the spectrum in HBM) or fp16 twiddles (breaking the 1e-4 contract).

Tuning

Per-n_fft frames-per-block defaults live in flashmel/transform.py (_FPB_DEFAULTS), found with scripts/bench.py --sweep. Override per process with FLASHMEL_FPB.

Development

uv sync --extra dev                          # flashmel deps + pytest
uv pip install torch torchaudio --index-url https://download.pytorch.org/whl/cu130
uv run --no-sync pytest scripts/test_correctness.py   # 87 tests vs torchaudio (needs CUDA)
uv run --no-sync python scripts/bench.py [--sweep]

torch/torchaudio aren't in the lockfile (see above), so install them into the venv separately and pass --no-sync to uv run so it doesn't prune them.

Correctness thresholds per PROJECT.md: atol=rtol=1e-4 (fp32 out), 1e-2 (fp16 out, compared against the fp16-quantized reference; values beyond fp16 range legitimately saturate to inf on both sides).

Layout: flashmel/ (package: transform.py wrapper, filters.py host-side window/filterbank packing, build.py NVRTC + cache, runtime.py driver-API launch, kernel.cu), scripts/ (tests + bench). The cubin cache lives in ~/.cache/flashmel (FLASHMEL_CACHE_DIR to override). cuFFTDx headers come from the nvidia-mathdx wheel.

License

MIT — see LICENSE.