Metadata-Version: 2.5
Name: flashmel
Version: 0.1.0
Summary: Fused CUDA mel spectrogram (cuFFTDx) — a fast, torchaudio-compatible MelSpectrogram in a single kernel launch.
Project-URL: Homepage, https://github.com/AlumKal/flashmel
Project-URL: Repository, https://github.com/AlumKal/flashmel
Project-URL: Issues, https://github.com/AlumKal/flashmel/issues
Author-email: AlumKal <alumkal-pub@outlook.com>
License-Expression: MIT
License-File: LICENSE
Keywords: audio,cuda,cufftdx,gpu,mel,spectrogram,stft,torchaudio
Classifier: Development Status :: 4 - Beta
Classifier: Environment :: GPU :: NVIDIA CUDA
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: POSIX :: Linux
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Multimedia :: Sound/Audio :: Analysis
Classifier: Topic :: Scientific/Engineering
Requires-Python: >=3.10
Requires-Dist: cuda-python<14,>=13.0
Requires-Dist: numpy>=1.24
Requires-Dist: nvidia-cuda-cccl<14,>=13.0
Requires-Dist: nvidia-cuda-nvrtc<14,>=13.0
Requires-Dist: nvidia-cuda-runtime<14,>=13.0
Requires-Dist: nvidia-mathdx>=25.6.0
Provides-Extra: dev
Requires-Dist: pytest>=9.0.3; extra == 'dev'
Description-Content-Type: text/markdown

# 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.:

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

```sh
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

```python
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_mode`s) 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):

| case | config | batch | ours | torchaudio | speedup | essential BW |
|---|---|---|---|---|---|---|
| whisper | 400/160, 128 mel, 30 s | 32 | 0.54 ms | 6.34 ms | 11.7× | 205 GB/s (103%) |
| whisper i16→fp16 | 400/160, 128 mel | 32 | 0.53 ms | 7.28 ms | 13.8× | 105 GB/s |
| speech | 1024/256, 80 mel | 32 | 0.20 ms | 3.15 ms | 16.2× | 138 GB/s (70%) |
| small | 64/16, 40 mel | 64 | 0.051 ms | 0.29 ms | 5.6× | 280 GB/s (141%) |
| music | 4096/1024, 256 mel, 44.1k | 16 | 0.35 ms | 4.54 ms | 12.9× | 100 GB/s (51%) |
| huge | 16384/4096, 256 mel, 44.1k | 16 | 0.92 ms | 4.61 ms | 5.0× | 33 GB/s |

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

```sh
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](LICENSE).
