A torch-native CWT module, a real GPU benchmark, and the bottleneck it uncovered

2026-08-20 · performance, wavelet coherence, GPU

The wavelet transform underneath every coherence computation in this project has always run through fcwt, an FFTW-backed C++ library called from a plain, unbatched Python loop — one (sample, channel) pair at a time, CPU-only, no GPU involvement at all. This session replaced it with a torch-native (torch.fft) implementation, wired it in as the real default, and along the way found and fixed a genuine bottleneck that had nothing to do with the CWT itself.

Building the module, and two real calibration bugs

The new module (utils/torch_cwt.py) computes the CWT via the Fourier convolution theorem — FFT the signal once, multiply against a precomputed frequency-domain Morlet filter bank, inverse-FFT back. Two real bugs turned up while calibrating against actual fcwt output rather than trusting the math in isolation: fcwt's frequency array runs high-to-low, the opposite of the natural construction order, which silently compared the wrong frequency's column against the wrong column until caught (magnitude correlation was ~0.22 until fixed, ~1.0 after); and the filter's peak amplitude needed an empirically-calibrated flat constant across scale, since a literal derivation gives an amplitude that wrongly grows at low frequency. That second bug alone accounted for a worst-case coefficient error of ~53 inside the region that actually reaches the model — high per-frequency correlation had been masking it, since correlation doesn't catch a systematic amplitude bias.

Parity validation

Magnitude and phase correlation between the new module and real fcwt output on actual CHB-MIT trials, restricted to the cone-of-influence region that's the only part that ever reaches the model: 0.999950–0.999962 magnitude Pearson r. A broader sweep across 60 trials, every channel, put the per-trial minimum at 0.99989. A stress test at 500 log-spaced scales on a real 300-second signal held the same correctness (median relative error 0.005%, no NaN/Inf across 38 million valid samples).

The real GPU benchmark

Local Mac testing (CPU and Apple's MPS backend) gave a mixed picture — a naive torch implementation sometimes lost to fcwt's own algorithmic edge on long signals. That didn't hold up on real hardware. A Runpod GPU pod (RTX 4090, ~8 minutes of uptime, about five cents) gave a clean answer instead:

Correctness held throughout — magnitude/phase correlation 0.999998–1.000000, no NaN or Inf at any size tested.

Wiring it in for real: batching, not a naive swap

Pointing the classifier at the new module unchanged would have erased most of the measured speedup — every real call site invoked the transform one single-channel signal at a time, and that per-item host↔device transfer and kernel launch is exactly the overhead the batched benchmark above didn't pay. So the call sites themselves were reworked: a batched transform function, a vectorized sibling of the existing feature-preparation code, and cache-key resolution done in a first pass so only actual cache misses get batched through the transform. A new cwt_backend flag ("fcwt" default, "torch" opt-in) keeps the old path byte-for-byte available as a revert switch rather than requiring a git revert.

Verified bit-for-bit identical output between the looped and batched call paths on synthetic data, then end-to-end on a trained classifier (predict_proba max absolute difference 0.0023, 100% class agreement between backends) — consistent with the CWT-level correlation propagating cleanly through training, not a bug.

A cache-key bug that silently no-op'd the first real comparison

The first real-data before/after comparison between backends came back with bit-identical scores for both — which turned out to be wrong, for a specific, findable reason. The disk caches for CWT and dense-edge features were keyed on signal content and config, but not on which backend actually computed the entry. This machine already had a populated cache from earlier fcwt-only sessions, so the "torch" backend's run was silently reading back fcwt-computed values on every lookup and never actually exercising the new module at all — both runs' progress bars showing 100% cache reuse is what gave it away. Fixed by folding the backend name into the cache key, which forces a one-time, harmless full recompute on the next run anywhere a cache already existed. With that fixed, the real comparison showed small, non-systematic differences in the continuous-score metrics per fold — consistent with the already-established near-perfect coefficient-level correlation, not a bug.

Default flipped, fcwt removed — and the real bottleneck surfaces

With the swap validated, cwt_backend="torch" became the real default for every pipeline, fcwt and its unused FFTW dependency were dropped entirely, and a pod image was built with the new, simpler dependency set baked in. But a real, full-scale attempt at the actual target evaluation (30s windows, prediction mode, full chb01) told a more important story: the CWT step itself was genuinely fast and batched, exactly as benchmarked — but a downstream step, dense-edge coherence computation, took about as long per batch regardless of how fast the CWT feeding it ran. Projected total time for the full evaluation: roughly 20 hours, essentially unchanged from the estimate before any of this session's work. The run was killed once that was clear, not because anything crashed.

Root cause: not the GPU math, the compressed cache write

The dense-edge computation itself — genuinely vectorized torch, no hidden Python loops — was confirmed correct by reading it in full before measuring anything. The actual cost turned out to be one line: every trial's computed dense-edge tensor was being written to disk with np.savez_compressed, synchronous single-threaded DEFLATE compression, timed inside the same loop as the GPU compute. Measured directly on a fresh pod with GPU-only timing isolated via torch.cuda.synchronize():

With the disk-write cost roughly 14× cheaper and GPU compute already fast, dense-edge should no longer dominate — but that's a projection from a 24-window slice, not a re-measurement of the full evaluation. The full run that was killed earlier hasn't been re-attempted yet, left for a future session rather than run overnight unmeasured.

Open items

The full real evaluation hasn't been re-run end-to-end since the cache-write fix, so there's no fresh measured wall-clock number yet, only the projection above. The cache-key fix means every pre-existing on-disk CWT/dense-edge cache entry (this machine and any Runpod volume that accumulated one earlier) gets invalidated once, forcing a one-time full recompute on next use — expected, not a regression. A demo script and its output image exist but aren't committed. run_pipelines.py's window-length and frequency-resolution defaults still reflect the outdated canonical config, not yet updated to match the real 30s-window target.