flashmel 0.1.0__tar.gz

flashmel-0.1.0/.gitignore

@@ -0,0 +1,6 @@
+__pycache__/
+.venv/
+uv.lock
+build/
+dist/
+.pytest_cache/

flashmel-0.1.0/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.

flashmel-0.1.0/PKG-INFO

@@ -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/PROJECT.md

@@ -0,0 +1,18 @@
+Your goal is to create a fused mel spectrogram kernel in cuFFTDx, and provide a python wrapper.
+
+You must support all arguments of <https://docs.pytorch.org/audio/stable/generated/torchaudio.transforms.MelSpectrogram.html>, except:
+
+- `window_fn` and `wkwargs`: always use the default
+- `n_fft`: assume the value is either 400 or power-of-2 between 2^6 and 2^14
+- `n_mels`: assume a maximum of 256
+
+And additionally support:
+
+- input dtype: i16 (raw pcm) or f32
+- output dtype: fp16 or fp32
+
+Write proper scripts for correctness test and timing before writing the kernel. Use `atol=1e-4, rtol=1e-4` (100x for fp16) as correctness threshold and cuda events for timing.
+
+Keep optimizing until your kernel saturates on memory bandwidth (likely) or compute.
+
+The project is empty now - decide a good directory structure yourself. See ../mel-test for a relevant example.

flashmel-0.1.0/README.md

@@ -0,0 +1,167 @@
+# 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/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-0.1.0/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-0.1.0/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