flashmel 0.1.0__py3-none-any.whl

flashmel/__init__.py

@@ -0,0 +1,12 @@
+"""flashmel: fused CUDA mel spectrogram (cuFFTDx), torchaudio-compatible."""
+
+from importlib.metadata import PackageNotFoundError, version
+
+from .transform import MelSpectrogram
+
+try:
+    __version__ = version("flashmel")
+except PackageNotFoundError:  # not installed (e.g. running from a source tree)
+    __version__ = "0.0.0+unknown"
+
+__all__ = ["MelSpectrogram", "__version__"]

flashmel/build.py

@@ -0,0 +1,80 @@
+"""NVRTC compilation of kernel.cu -> cubin, cached on disk by content hash."""
+
+import hashlib
+import importlib.util
+import os
+from pathlib import Path
+
+from cuda.bindings import nvrtc
+
+_KERNEL_SRC = Path(__file__).with_name("kernel.cu")
+
+
+def _cache_dir() -> Path:
+    """Per-user, env-overridable cubin cache (never inside site-packages)."""
+    if env := os.environ.get("FLASHMEL_CACHE_DIR"):
+        return Path(env)
+    base = os.environ.get("XDG_CACHE_HOME") or (Path.home() / ".cache")
+    return Path(base) / "flashmel"
+
+
+def _nvidia_root() -> Path:
+    # "nvidia" is a namespace package: pick the location that actually has mathdx
+    # (stray "nvidia" directories elsewhere on sys.path also join the namespace).
+    spec = importlib.util.find_spec("nvidia")
+    locations = [Path(p) for p in spec.submodule_search_locations]
+    for p in locations:
+        if (p / "mathdx" / "include").is_dir():
+            return p
+    raise RuntimeError(f"nvidia-mathdx include dir not found in {locations}")
+
+
+def include_dirs() -> list[Path]:
+    nv = _nvidia_root()
+    return [
+        nv / "mathdx" / "include",
+        nv / "mathdx" / "external" / "cutlass" / "include",
+        nv / "cu13" / "include",
+        nv / "cu13" / "include" / "cccl",
+    ]
+
+
+def _check(ret):
+    err, *rest = ret
+    if err != nvrtc.nvrtcResult.NVRTC_SUCCESS:
+        raise RuntimeError(f"NVRTC error: {err}")
+    return rest[0] if len(rest) == 1 else rest
+
+
+def compile_kernel(macros: dict[str, str], arch: str) -> bytes:
+    """Compile kernel.cu with the given -D macros for sm_{arch}; returns cubin bytes."""
+    src = _KERNEL_SRC.read_text()
+    opts = [
+        "--std=c++20",
+        f"--gpu-architecture=sm_{arch}",
+        "--device-as-default-execution-space",
+        "--use_fast_math",
+        *(f"-D{k}={v}" for k, v in sorted(macros.items())),
+        *(f"--include-path={d}" for d in include_dirs()),
+    ]
+    key = hashlib.sha256("\0".join([src, *opts]).encode()).hexdigest()[:16]
+    cache_dir = _cache_dir()
+    cached = cache_dir / f"flashmel_{key}.cubin"
+    if cached.exists():
+        return cached.read_bytes()
+
+    prog = _check(nvrtc.nvrtcCreateProgram(src.encode(), b"kernel.cu", 0, [], []))
+    try:
+        res = nvrtc.nvrtcCompileProgram(prog, len(opts), [o.encode() for o in opts])[0]
+        log = bytearray(_check(nvrtc.nvrtcGetProgramLogSize(prog)))
+        _check(nvrtc.nvrtcGetProgramLog(prog, log))
+        if res != nvrtc.nvrtcResult.NVRTC_SUCCESS:
+            raise RuntimeError(f"NVRTC compilation failed:\n{log.decode()}")
+        cubin = bytearray(_check(nvrtc.nvrtcGetCUBINSize(prog)))
+        _check(nvrtc.nvrtcGetCUBIN(prog, cubin))
+    finally:
+        nvrtc.nvrtcDestroyProgram(prog)
+
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    cached.write_bytes(bytes(cubin))
+    return bytes(cubin)

flashmel/filters.py

@@ -0,0 +1,48 @@
+"""Window construction and sparse mel filterbank packing (CPU, fp32)."""
+
+import torch
+from torchaudio.functional import melscale_fbanks
+
+
+def make_window(n_fft: int, win_length: int, normalized) -> torch.Tensor:
+    """Hann window with STFT normalization folded in, zero-padded centered to n_fft."""
+    w = torch.hann_window(win_length, periodic=True, dtype=torch.float32)
+    if normalized is True or normalized == "window":
+        w = w / w.pow(2).sum().sqrt()
+    elif normalized == "frame_length":
+        w = w / win_length**0.5
+    if win_length < n_fft:
+        left = (n_fft - win_length) // 2
+        w = torch.nn.functional.pad(w, (left, n_fft - win_length - left))
+    return w.contiguous()
+
+
+def pack_fbank(
+    n_freqs: int,
+    f_min: float,
+    f_max: float,
+    n_mels: int,
+    sample_rate: int,
+    norm,
+    mel_scale: str,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
+    """Pack the (triangular => contiguous-support) mel filterbank as sparse rows.
+
+    Returns (start [n_mels] i32, width [n_mels] i32, weights [n_mels, w_max] f32,
+    w_max) with start[m] clamped so start[m] + w_max <= n_freqs; weights are the
+    exact melscale_fbanks values (zero outside each filter's support).
+    """
+    fb = melscale_fbanks(n_freqs, f_min, f_max, n_mels, sample_rate, norm, mel_scale)
+    fb = fb.T.contiguous().to(torch.float32)  # (n_mels, n_freqs)
+    start = torch.zeros(n_mels, dtype=torch.int64)
+    width = torch.zeros(n_mels, dtype=torch.int64)
+    for m in range(n_mels):
+        nz = fb[m].nonzero().flatten()
+        if nz.numel():
+            start[m] = nz[0]
+            width[m] = nz[-1] - nz[0] + 1
+    w_max = max(int(width.max()), 1)
+    clamped = start.clamp(0, n_freqs - w_max)
+    width = width + start - clamped  # support end relative to the clamped start
+    weights = torch.stack([fb[m, clamped[m] : clamped[m] + w_max] for m in range(n_mels)])
+    return clamped.to(torch.int32).contiguous(), width.to(torch.int32).contiguous(), weights.contiguous(), w_max

flashmel/kernel.cu

@@ -0,0 +1,266 @@
+// Fused torchaudio-compatible MelSpectrogram kernel:
+//   framing + signal padding + Hann window + R2C FFT (cuFFTDx) + |X|^power
+//   + sparse mel filterbank projection, all in one launch.
+//
+// Compile-time macros injected by build.py:
+//   N_FFT       400 or power of 2 in [64, 16384]
+//   FPB         frames (FFTs) per block
+//   W_MAX       max mel filter support width
+//   N_MELS      number of mel bins
+//   CENTER      0/1 torch.stft center
+//   PAD_MODE    0 reflect, 1 constant(zeros), 2 replicate, 3 circular
+//   POWER_MODE  0: |X|^2, 1: |X|, 2: |X|^POWER_VAL (generic)
+//   POWER_VAL   float literal, e.g. 0.7f (only used when POWER_MODE == 2)
+//   IN_I16      0: f32 input, 1: i16 raw PCM input (scaled by 1/32768)
+//   OUT_HALF    0: f32 output, 1: fp16 output
+//   REAL_MODE   0: real_mode::normal (e.g. 400), 1: real_mode::folded (powers of 2)
+//   SM_ARCH     e.g. 890
+
+#include <cufftdx.hpp>
+#include <cuda_fp16.h>
+
+using namespace cufftdx;
+
+constexpr unsigned N_FREQS = N_FFT / 2 + 1;
+// Power-tile row stride: odd for every supported n_fft => rows hit distinct banks.
+constexpr unsigned PSTRIDE = N_FREQS;
+static_assert(PSTRIDE % 2 == 1);
+
+#if IN_I16
+using in_t = short;
+#else
+using in_t = float;
+#endif
+#if OUT_HALF
+using out_t = __half;
+#else
+using out_t = float;
+#endif
+
+#if REAL_MODE == 1
+using fft_real_opts = RealFFTOptions<complex_layout::natural, real_mode::folded>;
+#else
+using fft_real_opts = RealFFTOptions<complex_layout::natural, real_mode::normal>;
+#endif
+using FFT = decltype(Size<N_FFT>() + Precision<float>() + Type<fft_type::r2c>()
+                     + fft_real_opts() + Direction<fft_direction::forward>()
+                     + FFTsPerBlock<FPB>() + Block() + SM<SM_ARCH>());
+
+using complex_type = typename FFT::value_type;
+
+constexpr unsigned NTHR = FFT::block_dim.x * FFT::block_dim.y;
+
+// Mel projection lanes per output: when the block has more threads than mel
+// outputs (large n_fft => FPB=1), SPLIT threads cooperate on one dot product
+// and reduce via warp shuffle instead of idling.
+constexpr unsigned OUTPUTS = N_MELS * FPB;
+constexpr unsigned SPLIT = []() consteval {
+    // Narrow filters have nothing to split; partial warps would need shfl masks
+    // excluding non-existent lanes.
+    if (W_MAX <= 16 || NTHR % 32 != 0) return 1u;
+    unsigned s = 1;
+    while (s < 32 && 2 * s * OUTPUTS <= NTHR) s *= 2;
+    return s;
+}();
+constexpr unsigned SMEM_BYTES =
+    FFT::shared_memory_size > FPB * PSTRIDE * 4 ? FFT::shared_memory_size : FPB * PSTRIDE * 4;
+
+// Fetch sample v of the virtually padded signal: the f32/i16 signal x[0..n) is
+// first zero-padded by `pad` on both sides (length L = n + 2*pad), then extended
+// by PAD_MODE for the centered STFT overhang. Python validates the torch.stft
+// constraints that make a single reflection/wrap sufficient.
+__device__ __forceinline__ float load_sample(const in_t* __restrict__ x, int v, int L, int n,
+                                             int pad) {
+#if PAD_MODE == 0
+    v = v < 0 ? -v : v;
+    v = v >= L ? 2 * L - 2 - v : v;
+#elif PAD_MODE == 1
+    if (v < 0 || v >= L) return 0.f;
+#elif PAD_MODE == 2
+    v = v < 0 ? 0 : v;
+    v = v >= L ? L - 1 : v;
+#else
+    v %= L;
+    v = v < 0 ? v + L : v;
+#endif
+    v -= pad;
+    if (v < 0 || v >= n) return 0.f;
+#if IN_I16
+    return (float)x[v] * (1.f / 32768.f);
+#else
+    return x[v];
+#endif
+}
+
+__device__ __forceinline__ float spec_value(complex_type v) {
+    const float p2 = v.real() * v.real() + v.imag() * v.imag();
+#if POWER_MODE == 0
+    return p2;
+#elif POWER_MODE == 1
+    return sqrtf(p2);
+#else
+    return powf(p2, 0.5f * POWER_VAL);
+#endif
+}
+
+extern "C" __global__ void flashmel(const in_t* __restrict__ in,        // [B, n_samples]
+                                    out_t* __restrict__ out,            // [B, N_MELS, n_frames]
+                                    const float* __restrict__ window,   // [N_FFT]
+                                    const int* __restrict__ start,      // [N_MELS]
+                                    const int* __restrict__ width,      // [N_MELS]
+                                    const float* __restrict__ weights,  // [N_MELS, W_MAX]
+                                    int n_samples, int pad, int hop, int n_frames) {
+    complex_type thread_data[FFT::storage_size];
+    extern __shared__ __align__(16) unsigned char smem_raw[];
+    float* power = reinterpret_cast<float*>(smem_raw);
+
+    const unsigned b = blockIdx.y;
+    const unsigned frame0 = blockIdx.x * FPB;
+    const int L = n_samples + 2 * pad;
+    const in_t* x = in + (size_t)b * n_samples;
+
+    // Framing + padding + window, straight into FFT registers.
+    const unsigned frame = frame0 + threadIdx.y;
+    const bool valid = frame < (unsigned)n_frames;
+    const int base = (int)frame * hop - (CENTER ? (int)(N_FFT / 2) : 0);
+#if REAL_MODE == 1
+    // Folded R2C: N_FFT reals enter as N_FFT/2 complex (even, odd) pairs.
+#if IN_I16
+    using in2_t = short2;
+#else
+    using in2_t = float2;
+#endif
+    // Interior frames (no padding in play, pointer pair-aligned — odd batch row
+    // lengths shift row bases by 4 bytes) take vector loads: one 2-wide load
+    // each for signal and window instead of two scalars.
+    const in_t* xf = x + (base - pad);
+    if (valid && base >= pad && base + (int)N_FFT <= pad + n_samples &&
+        (reinterpret_cast<uintptr_t>(xf) & (sizeof(in2_t) - 1)) == 0) {
+        const in2_t* xp = reinterpret_cast<const in2_t*>(xf);
+        const float2* wp = reinterpret_cast<const float2*>(window);
+#pragma unroll
+        for (unsigned i = 0; i < FFT::input_ept; ++i) {
+            const unsigned c = threadIdx.x + i * FFT::stride;
+            float re = 0.f, im = 0.f;
+            if (c < FFT::input_length) {
+                const in2_t s = xp[c];
+                const float2 wv = wp[c];
+#if IN_I16
+                re = (float)s.x * (1.f / 32768.f) * wv.x;
+                im = (float)s.y * (1.f / 32768.f) * wv.y;
+#else
+                re = s.x * wv.x;
+                im = s.y * wv.y;
+#endif
+            }
+            thread_data[i] = complex_type(re, im);
+        }
+    } else {
+#pragma unroll
+        for (unsigned i = 0; i < FFT::input_ept; ++i) {
+            const unsigned c = threadIdx.x + i * FFT::stride;
+            float re = 0.f, im = 0.f;
+            if (valid && c < FFT::input_length) {
+                re = load_sample(x, base + (int)(2 * c), L, n_samples, pad) * window[2 * c];
+                im = load_sample(x, base + (int)(2 * c + 1), L, n_samples, pad) * window[2 * c + 1];
+            }
+            thread_data[i] = complex_type(re, im);
+        }
+    }
+#else
+    auto* reg = reinterpret_cast<float*>(thread_data);
+    // Interior frames (no padding in play) skip the per-sample boundary logic.
+    const in_t* xf = x + (base - pad);
+    if (valid && base >= pad && base + (int)N_FFT <= pad + n_samples) {
+#pragma unroll
+        for (unsigned i = 0; i < FFT::input_ept; ++i) {
+            const unsigned idx = threadIdx.x + i * FFT::stride;
+            float s = 0.f;
+            if (idx < FFT::input_length) {
+#if IN_I16
+                s = (float)xf[idx] * (1.f / 32768.f) * window[idx];
+#else
+                s = xf[idx] * window[idx];
+#endif
+            }
+            reg[i] = s;
+        }
+    } else {
+#pragma unroll
+        for (unsigned i = 0; i < FFT::input_ept; ++i) {
+            const unsigned idx = threadIdx.x + i * FFT::stride;
+            reg[i] = (valid && idx < FFT::input_length)
+                         ? load_sample(x, base + (int)idx, L, n_samples, pad) * window[idx]
+                         : 0.f;
+        }
+    }
+#endif
+
+    FFT().execute(thread_data, reinterpret_cast<complex_type*>(smem_raw));
+    __syncthreads();
+
+    // FFT workspace is dead; reuse it as the power tile.
+    float* prow = power + threadIdx.y * PSTRIDE;
+#pragma unroll
+    for (unsigned i = 0; i < FFT::output_ept; ++i) {
+        const unsigned idx = threadIdx.x + i * FFT::stride;
+        if (idx < N_FREQS) prow[idx] = spec_value(thread_data[i]);
+    }
+    __syncthreads();
+
+    // Sparse mel projection; consecutive threads take consecutive frames of the
+    // same mel row -> coalesced FPB-wide segments in out.
+    const unsigned tid = threadIdx.y * FFT::block_dim.x + threadIdx.x;
+    out_t* o = out + (size_t)b * N_MELS * n_frames + frame0;
+    if constexpr (SPLIT > 1) {
+        const unsigned lane = tid & (SPLIT - 1);
+        const unsigned group = tid / SPLIT;
+        constexpr unsigned NGROUPS = NTHR / SPLIT;
+#pragma unroll 1
+        for (unsigned pb = 0; pb < OUTPUTS; pb += NGROUPS) {  // uniform: all lanes shuffle
+            const unsigned p = pb + group;
+            const unsigned m = p / FPB, f = p % FPB;
+            const bool active = p < OUTPUTS && frame0 + f < (unsigned)n_frames;
+            float acc = 0.f;
+            if (active) {
+                const float* w = weights + (size_t)m * W_MAX;
+                const float* pr = power + f * PSTRIDE + start[m];
+                const int wm = width[m];
+#pragma unroll 4
+                for (int j = (int)lane; j < wm; j += SPLIT) acc = fmaf(w[j], pr[j], acc);
+            }
+#pragma unroll
+            for (unsigned s = SPLIT / 2; s > 0; s >>= 1)
+                acc += __shfl_down_sync(0xffffffffu, acc, s, SPLIT);
+            if (active && lane == 0) o[(size_t)m * n_frames + f] = (out_t)acc;
+        }
+    } else {
+#pragma unroll 1
+        for (unsigned p = tid; p < OUTPUTS; p += NTHR) {
+            const unsigned m = p / FPB, f = p % FPB;
+            if (frame0 + f >= (unsigned)n_frames) continue;
+            const float* w = weights + (size_t)m * W_MAX;
+            const float* pr = power + f * PSTRIDE + start[m];
+            float acc = 0.f;
+#if W_MAX <= 16
+            // Narrow filters: fixed fully-unrolled loop over the zero-padded row.
+#pragma unroll
+            for (int j = 0; j < W_MAX; ++j) acc = fmaf(w[j], pr[j], acc);
+#else
+            // Wide filters: most are much narrower than W_MAX; looping to the actual
+            // support width cuts shared-memory traffic (the L1/SMEM bottleneck).
+            const int wm = width[m];
+#pragma unroll 4
+            for (int j = 0; j < wm; ++j) acc = fmaf(w[j], pr[j], acc);
+#endif
+            o[(size_t)m * n_frames + f] = (out_t)acc;
+        }
+    }
+}
+
+extern "C" __global__ void flashmel_info(unsigned* o) {
+    o[0] = FFT::block_dim.x;
+    o[1] = FFT::block_dim.y;
+    o[2] = SMEM_BYTES;
+    o[3] = FPB;
+}

flashmel/runtime.py

@@ -0,0 +1,65 @@
+"""CUDA driver-API module loading and kernel launches on torch's current stream."""
+
+import functools
+
+import numpy as np
+import torch
+from cuda.bindings import driver
+
+from .build import compile_kernel
+
+
+def _check(ret):
+    err, *rest = ret
+    if err != driver.CUresult.CUDA_SUCCESS:
+        raise RuntimeError(f"CUDA driver error: {err}")
+    return rest[0] if len(rest) == 1 else rest
+
+
+def _launch_raw(fn, grid, block, smem, stream, *args):
+    vals = np.array(args, dtype=np.uint64)
+    ptrs = np.array([vals.ctypes.data + 8 * i for i in range(len(args))], dtype=np.uint64)
+    _check(driver.cuLaunchKernel(fn, *grid, *block, smem, stream, ptrs.ctypes.data, 0))
+
+
+class KernelModule:
+    """One compiled (cubin) flashmel kernel for a fixed macro configuration.
+
+    A CUmodule is bound to the CUDA context it was loaded in, so each instance
+    is tied to one device: it loads under (and launches on) torch's primary
+    context for `device_index`.
+    """
+
+    def __init__(self, macros: dict[str, str], arch: str, device_index: int) -> None:
+        cubin = compile_kernel(macros, arch)
+        self.device_index = device_index
+        with torch.cuda.device(device_index):
+            self.mod = _check(driver.cuModuleLoadData(cubin))
+            self.fn = _check(driver.cuModuleGetFunction(self.mod, b"flashmel"))
+            info_fn = _check(driver.cuModuleGetFunction(self.mod, b"flashmel_info"))
+
+            info = torch.zeros(4, dtype=torch.uint32, device=f"cuda:{device_index}")
+            stream = torch.cuda.current_stream().cuda_stream
+            _launch_raw(info_fn, (1, 1, 1), (1, 1, 1), 0, stream, info.data_ptr())
+            torch.cuda.synchronize()
+            bx, by, self.smem, self.fpb = info.cpu().tolist()
+            self.block = (bx, by, 1)
+            if self.smem > 48 * 1024:
+                _check(
+                    driver.cuFuncSetAttribute(
+                        self.fn,
+                        driver.CUfunction_attribute.CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
+                        self.smem,
+                    )
+                )
+
+    def launch(self, n_frames: int, batch: int, *args) -> None:
+        with torch.cuda.device(self.device_index):
+            stream = torch.cuda.current_stream().cuda_stream
+            grid = ((n_frames + self.fpb - 1) // self.fpb, batch, 1)
+            _launch_raw(self.fn, grid, self.block, self.smem, stream, *args)
+
+
+@functools.lru_cache(maxsize=None)
+def get_module(macro_items: tuple[tuple[str, str], ...], arch: str, device_index: int) -> KernelModule:
+    return KernelModule(dict(macro_items), arch, device_index)

flashmel/transform.py

@@ -0,0 +1,268 @@
+"""MelSpectrogram transform mirroring torchaudio.transforms.MelSpectrogram."""
+
+import math
+import os
+
+import torch
+
+from .filters import make_window, pack_fbank
+from .runtime import KernelModule, get_module
+
+_VALID_N_FFT = {400} | {1 << k for k in range(6, 15)}
+_PAD_MODES = {"reflect": 0, "constant": 1, "replicate": 2, "circular": 3}
+_MAX_N_MELS = 256
+_MAX_GRID_Y = 65535
+
+# Per-n_fft tuning defaults from scripts/bench.py --sweep on RTX 4070 Laptop
+# (overridable via FLASHMEL_FPB).
+_FPB_DEFAULTS = {64: 32, 128: 32, 256: 16, 400: 8, 512: 8, 1024: 4, 2048: 4, 4096: 2, 8192: 1, 16384: 1}
+
+
+def _resolve_module(
+    n_fft: int,
+    w_max: int,
+    n_mels: int,
+    center: bool,
+    pad_mode: str,
+    power: float,
+    in_i16: bool,
+    out_half: bool,
+    device: torch.device,
+) -> KernelModule:
+    if power == 2.0:
+        power_mode = 0
+    elif power == 1.0:
+        power_mode = 1
+    else:
+        power_mode = 2
+    major, minor = torch.cuda.get_device_capability(device)
+    arch = f"{major}{minor}"
+    fpb = int(os.environ.get("FLASHMEL_FPB", _FPB_DEFAULTS[n_fft]))
+    last_err: Exception | None = None
+    while fpb >= 1:
+        macros = {
+            "N_FFT": str(n_fft),
+            "FPB": f"{fpb}u",
+            "W_MAX": str(w_max),
+            "N_MELS": str(n_mels),
+            "CENTER": "1" if center else "0",
+            "PAD_MODE": str(_PAD_MODES[pad_mode]),
+            "POWER_MODE": str(power_mode),
+            "POWER_VAL": f"{power!r}f",
+            "IN_I16": "1" if in_i16 else "0",
+            "OUT_HALF": "1" if out_half else "0",
+            # powers of 2: cuFFTDx folded R2C; 400: normal-mode R2C (see README for
+            # the rejected paired-c2c alternative).
+            "REAL_MODE": "1" if n_fft != 400 else "0",
+            "SM_ARCH": f"{arch}0",
+        }
+        try:
+            return get_module(tuple(sorted(macros.items())), arch, device.index)
+        except RuntimeError as e:  # invalid cuFFTDx config (e.g. too many threads)
+            last_err = e
+            fpb //= 2
+    raise RuntimeError(f"no valid kernel config found for n_fft={n_fft}") from last_err
+
+
+def _n_frames(n: int, n_fft: int, hop_length: int, pad: int, center: bool) -> int:
+    L = n + 2 * pad
+    return L // hop_length + 1 if center else (L - n_fft) // hop_length + 1
+
+
+# Opaque custom op so torch.compile traces through the transform without graph
+# breaks on the driver-API launch (shapes come from the register_fake impl).
+@torch.library.custom_op("flashmel::mel_forward", mutates_args=(), device_types="cuda")
+def _mel_forward(
+    x: torch.Tensor,
+    window: torch.Tensor,
+    fb_start: torch.Tensor,
+    fb_width: torch.Tensor,
+    fb_weights: torch.Tensor,
+    n_fft: int,
+    hop_length: int,
+    pad: int,
+    center: bool,
+    pad_mode: str,
+    power: float,
+    out_half: bool,
+) -> torch.Tensor:
+    x = x.contiguous()
+    batch, n = x.shape
+    n_mels, w_max = fb_weights.shape
+    n_frames = _n_frames(n, n_fft, hop_length, pad, center)
+    out_dtype = torch.float16 if out_half else torch.float32
+    out = torch.empty(batch, n_mels, n_frames, device=x.device, dtype=out_dtype)
+    if out.numel():
+        mod = _resolve_module(
+            n_fft, w_max, n_mels, center, pad_mode, power, x.dtype == torch.int16, out_half, x.device
+        )
+        for b0 in range(0, batch, _MAX_GRID_Y):
+            b1 = min(b0 + _MAX_GRID_Y, batch)
+            mod.launch(
+                n_frames,
+                b1 - b0,
+                x[b0:b1].data_ptr(),
+                out[b0:b1].data_ptr(),
+                window.data_ptr(),
+                fb_start.data_ptr(),
+                fb_width.data_ptr(),
+                fb_weights.data_ptr(),
+                n,
+                pad,
+                hop_length,
+                n_frames,
+            )
+    return out
+
+
+@_mel_forward.register_fake
+def _mel_forward_fake(x, window, fb_start, fb_width, fb_weights, n_fft, hop_length, pad, center, pad_mode, power, out_half):
+    n_frames = _n_frames(x.shape[-1], n_fft, hop_length, pad, center)
+    out_dtype = torch.float16 if out_half else torch.float32
+    return x.new_empty((x.shape[0], fb_weights.shape[0], n_frames), dtype=out_dtype)
+
+
+class MelSpectrogram(torch.nn.Module):
+    """Drop-in fused-CUDA replacement for torchaudio.transforms.MelSpectrogram.
+
+    Differences from torchaudio:
+      - window_fn/wkwargs are not accepted (always periodic Hann);
+      - n_fft must be 400 or a power of 2 in [64, 16384]; n_mels <= 256;
+      - input must be a CUDA tensor of dtype float32 or int16 (raw PCM, scaled
+        by 1/32768);
+      - `dtype` selects the output dtype: torch.float32 (default) or torch.float16;
+      - inference-only: raises if the input requires grad while grad mode is on.
+    """
+
+    def __init__(
+        self,
+        sample_rate: int = 16000,
+        n_fft: int = 400,
+        win_length: int | None = None,
+        hop_length: int | None = None,
+        f_min: float = 0.0,
+        f_max: float | None = None,
+        pad: int = 0,
+        n_mels: int = 128,
+        power: float = 2.0,
+        normalized: bool | str = False,
+        center: bool = True,
+        pad_mode: str = "reflect",
+        onesided: bool = True,
+        norm: str | None = None,
+        mel_scale: str = "htk",
+        dtype: torch.dtype = torch.float32,
+    ) -> None:
+        super().__init__()
+        if n_fft not in _VALID_N_FFT:
+            raise ValueError(f"n_fft must be 400 or a power of 2 in [64, 16384], got {n_fft}")
+        win_length = win_length if win_length is not None else n_fft
+        hop_length = hop_length if hop_length is not None else win_length // 2
+        if not 0 < win_length <= n_fft:
+            raise ValueError(f"expected 0 < win_length <= n_fft, got {win_length}")
+        if hop_length <= 0:
+            raise ValueError(f"hop_length must be positive, got {hop_length}")
+        if not 1 <= n_mels <= _MAX_N_MELS:
+            raise ValueError(f"n_mels must be in [1, {_MAX_N_MELS}], got {n_mels}")
+        if pad < 0:
+            raise ValueError(f"pad must be non-negative, got {pad}")
+        power = float(power)
+        if not power > 0 or not math.isfinite(power):
+            raise ValueError(f"power must be a positive finite float, got {power}")
+        if normalized not in (True, False, "window", "frame_length"):
+            raise ValueError(f"normalized must be bool, 'window' or 'frame_length', got {normalized!r}")
+        if pad_mode not in _PAD_MODES:
+            raise ValueError(f"pad_mode must be one of {sorted(_PAD_MODES)}, got {pad_mode!r}")
+        if not onesided:
+            raise ValueError("onesided=False is not supported (mel scale needs n_fft // 2 + 1 bins)")
+        if norm is not None and norm != "slaney":
+            raise ValueError(f"norm must be None or 'slaney', got {norm!r}")
+        if mel_scale not in ("htk", "slaney"):
+            raise ValueError(f"mel_scale must be 'htk' or 'slaney', got {mel_scale!r}")
+        if dtype not in (torch.float32, torch.float16):
+            raise ValueError(f"dtype must be torch.float32 or torch.float16, got {dtype}")
+
+        self.sample_rate = sample_rate
+        self.n_fft = n_fft
+        self.win_length = win_length
+        self.hop_length = hop_length
+        self.f_min = f_min
+        self.f_max = f_max or float(sample_rate // 2)  # mirrors torchaudio exactly
+        self.pad = pad
+        self.n_mels = n_mels
+        self.power = power
+        self.normalized = normalized
+        self.center = center
+        self.pad_mode = pad_mode
+        self.norm = norm
+        self.mel_scale = mel_scale
+        self.dtype = dtype
+
+        n_freqs = n_fft // 2 + 1
+        start, width, weights, self.w_max = pack_fbank(
+            n_freqs, self.f_min, self.f_max, n_mels, sample_rate, norm, mel_scale
+        )
+        self.register_buffer("window", make_window(n_fft, win_length, normalized))
+        self.register_buffer("fb_start", start)
+        self.register_buffer("fb_width", width)
+        self.register_buffer("fb_weights", weights)
+        self._filters_by_device: dict[int, tuple[torch.Tensor, ...]] = {}
+
+    def _filters_on(self, device: torch.device) -> tuple[torch.Tensor, ...]:
+        """Window + filterbank tensors on `device`, cached per device (the
+        registered buffers stay wherever the module was moved to)."""
+        if self.window.device == device:
+            return self.window, self.fb_start, self.fb_width, self.fb_weights
+        cached = self._filters_by_device.get(device.index)
+        if cached is None:
+            cached = tuple(
+                t.to(device) for t in (self.window, self.fb_start, self.fb_width, self.fb_weights)
+            )
+            self._filters_by_device[device.index] = cached
+        return cached
+
+    def forward(self, waveform: torch.Tensor) -> torch.Tensor:
+        if waveform.dtype not in (torch.float32, torch.int16):
+            raise TypeError(f"input dtype must be float32 or int16, got {waveform.dtype}")
+        if not waveform.is_cuda:
+            raise ValueError("input must be a CUDA tensor")
+        if waveform.requires_grad and torch.is_grad_enabled():
+            raise RuntimeError(
+                "flashmel.MelSpectrogram is inference-only and does not support autograd; "
+                "run under torch.no_grad() or detach the input"
+            )
+
+        shape = waveform.shape
+        n = shape[-1]
+        L = n + 2 * self.pad
+        half = self.n_fft // 2
+        if self.center:
+            if self.pad_mode == "reflect" and half >= L:
+                raise ValueError(f"reflect padding requires input length > {half}, got {L}")
+            if self.pad_mode == "circular" and half > L:
+                raise ValueError(f"circular padding requires input length >= {half}, got {L}")
+            if L < 1:
+                raise ValueError("input is empty")
+            n_frames = L // self.hop_length + 1
+        else:
+            if L < self.n_fft:
+                raise ValueError(f"center=False requires input length >= n_fft ({self.n_fft}), got {L}")
+            n_frames = (L - self.n_fft) // self.hop_length + 1
+
+        x = waveform.reshape(-1, n).contiguous()
+        window, fb_start, fb_width, fb_weights = self._filters_on(waveform.device)
+        out = _mel_forward(
+            x,
+            window,
+            fb_start,
+            fb_width,
+            fb_weights,
+            self.n_fft,
+            self.hop_length,
+            self.pad,
+            self.center,
+            self.pad_mode,
+            self.power,
+            self.dtype == torch.float16,
+        )
+        return out.reshape(*shape[:-1], self.n_mels, n_frames)

flashmel-0.1.0.dist-info/METADATA

@@ -0,0 +1,201 @@
+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).

flashmel-0.1.0.dist-info/RECORD

@@ -0,0 +1,10 @@
+flashmel/__init__.py,sha256=G9wiMGmcgV1liMWuNUlt2zoNHdYg3ER4JIXOlLuL4Jg,381
+flashmel/build.py,sha256=iKKjOh6EcilxNZLtVjmbeKZJhaAMoWxFhKGLsBhSHwU,2823
+flashmel/filters.py,sha256=3A5xwI8qcfNGFutJSpdy6iMW7JZp7Vw5KQ33EQO3aNc,2013
+flashmel/kernel.cu,sha256=S0P6nkqh3lNv2ZCPBK2J4LoclgmFdw-iZ7kPQAqw16E,10390
+flashmel/runtime.py,sha256=DhEabDnfntRvxWS2QBRBmOPgDSF6HdLxyT0eBdInttA,2621
+flashmel/transform.py,sha256=C1l5ypZtrB18ZLvnGYafz_T4eAbEfxptOwDxO2ZTJpE,10613
+flashmel-0.1.0.dist-info/METADATA,sha256=pwmAD7VGv_Djco5MEwFJEruoLnoJquYf5QpFjIb0lUU,9872
+flashmel-0.1.0.dist-info/WHEEL,sha256=W3fkpkm7-wf9vBI5Z-7s0eWkeM-spu78I8Neb98DeEg,87
+flashmel-0.1.0.dist-info/licenses/LICENSE,sha256=abKhAw1lauLcmo-Lv5EUk_Hy9CK0OV0dG1QX_aVOkoY,1064
+flashmel-0.1.0.dist-info/RECORD,,

flashmel-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.32.4
+Root-Is-Purelib: true
+Tag: py3-none-any

flashmel-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 AlumKal
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.