Optimize SamplingDelayAlignment for MLX and right-size its fractional-delay FIR

[cboulay]

Closes #35.

MLX depthwise-conv FIR

The per-tap loop is a fine portable fallback, but on MLX it dispatches filter_len multiply-add stages per message — and at live chunk sizes dispatch, not arithmetic, is the cost. The FIR now runs as a single mx.conv_general with one group per column. The kernel (per-column taps, reversed, (n_cols, filter_len, 1)) is laid out once in _reset_state and cached in state, never rebuilt per message. Every other backend keeps the tap-sum unchanged.

Per-message ms, 256 channels, MLX synchronized, from the new benchmark:

samples/msg	conv (13 taps)	tap loop (13)	conv (33 taps)	tap loop (33)
3	0.137	0.350	0.181	0.575
30	0.145	0.333	0.155	0.589
300	0.172	0.381	0.151	0.472
1000	0.161	0.418	0.257	0.717
mixed	0.158	0.351	0.195	0.622

2.2×–3.8×, and roughly flat in chunk size where the loop was not.

filter_len 33 → 13

The requirement I sized against: over the intended passband, across every within-bank slot, at most 0.05° phase error and 0.01 dB magnitude error — about three orders of magnitude below the ~81° of skew the transformer exists to remove.

Passband	filter_len	Max phase err	Max mag err	Bulk delay
0–500 Hz (LFP)	7	0.0009°	0.00001 dB	3 samples
0–3 kHz	9	0.0038°	0.0001 dB	4 samples
0–7.5 kHz	11	0.091° ✗	0.014 dB	5 samples
0–7.5 kHz	13	0.015°	0.0015 dB	6 samples
0–7.5 kHz	33	0.0028°	0.0004 dB	16 samples

13 is the shortest odd length that holds the tolerance over the full broadband/spike band, so it becomes the default: latency drops from 16 samples to 6 (533 µs → 200 µs at 30 kHz) at error still far below the noise floor. test_filter_response_meets_tolerance pins each row against the filters the transformer actually designs (read back from state), and a companion test asserts 11 misses — so the default can't silently drift down. Recommended values by bandwidth are documented in the filter_len docstring and the processing guide.

This is the one user-visible behavior change in the PR: pipelines relying on the old 16-sample bulk delay will see 6 instead (the time-axis offset compensates, as before), and the output differs by the filters' small response difference. Set filter_len=33 to keep the old behavior exactly.

Rail handling

Profiling it separately, as the issue suggested, showed the Hillis-Steele forward-fill becomes the dominant per-message cost once the FIR is fast — 5.79 ms/message at 1000 samples × 256 ch on numpy, dwarfing the 0.91 ms FIR. The running max now uses the one-pass form where the backend has it: mx.cummax on MLX, maximum.accumulate on numpy/cupy. numpy's worst case drops to 3.32 ms; MLX rail handling now adds ~0.04 ms to a message. The log-depth scan stays for torch/jax and is still tested, via a namespace shim that hides maximum.accumulate — otherwise it would have lost coverage entirely, since torch isn't installed in CI.

Also

• Fixed history carry-over for filter_len=1: xext[-(n_taps-1):] is xext[-0:], which kept the entire extended buffer instead of nothing, so a 1-tap filter grew its history without bound. Now sliced from the length.
• examples/bench_sampling_delay_alignment.py — per-message timings across message sizes (including a mixed stream), channel counts, rail on/off, and conv-vs-tap-loop, with every MLX message evaluated before the clock is read.
• Relaxed the cross-backend comparison to an explicit float32 tolerance (rtol=1e-6, atol=1e-5). The conv accumulates in a different order than numpy's tap-sum; on the test's held-rail samples (~1e4) that is ~2e-3 absolute, i.e. 1.8e-7 relative — float32 eps, not a behavior difference.

Verification

112 passed, 1 skipped (torch not installed), ruff check and format clean. New tests cover the response table, the default's justification, conv-vs-tap-loop equivalence on MLX, chunk invariance through the conv path, kernel caching across messages, the scan/one-pass rail-fill agreement, and the filter_len=1 edge.

🤖 Generated with Claude Code

---

CI flake (second commit)

The first CI run failed fetching dnss256.zip with HTTP 500 from GitHub's release-asset CDN — unrelated to this change. The asset is fine (3/3 at HTTP 200 when probed), and a re-run of the same commit went green, so it was a blip. But all 12 matrix jobs fetch these fresh with no cache, so any single blip fails the build.

tests/conftest.py now retries the download up to 4 times with exponential backoff, skipping retries on 4xx so a genuinely missing asset still fails fast on the first attempt. Each attempt downloads to a .part file and renames into place, so an interrupted download can't leave a truncated zip that a later run treats as cached.

---

numba fast path for the numpy backend (follow-up commit)

Following review discussion on whether the numpy tap loop could do better with scipy: benchmarking showed scipy's fftconvolve loses to the tap loop at live chunk sizes and only wins for 33-tap offline batches — a regime we moved away from. numba, however, wins everywhere, so it (not scipy) becomes the numpy fast path, gated on an optional accel extra. Without numba, the portable Array-API path runs unchanged, so nothing is required; torch/jax/cupy keep the portable path regardless.

FIR — a jitted depthwise multiply-accumulate, channel loop innermost so it stays contiguous and auto-vectorizes, taps cached per column at state reset like the MLX kernel. ~3× faster than the tap loop at live sizes; a threaded variant engages above 4096 samples for offline batch work. Per-message ms, 256 ch, 13 taps (fir = numba, portable = tap loop):

samples/msg	numba fir	portable	speedup
3	0.006	0.019	2.9×
30	0.014	0.042	3.0×
300	0.077	0.235	3.1×
1000	0.254	0.816	3.2×
30 000	3.56	47.7	13.4×
300 000	37.6	700	18.6×

Rail forward-fill — the single left-to-right pass it actually is, replacing the O(n log n) scan. Columns are independent, so the parallel variant fans the per-column time recurrence across channels; this made it the fastest option at every size measured. A useful surprise from profiling: numpy's maximum.accumulate (the one-pass form added earlier in this PR) is actually the slowest of the three, dominated by the index gather — the numba parallel fill is ~7× faster than it at 300k samples.

Live-safety threshold — the 4096-sample split keeps live-sized messages on the serial kernels, so a live ezmsg graph never spawns a numba thread pool per message (which would add latency jitter and risk oversubscription with other numba users). Only offline batch buffers parallelize.

Packaging — kernels live in a separate _numba_kernels module so numba's import and first-call compile stay off the path of callers that never touch the numpy backend. numba>=0.59 is a new accel optional-dependency and is added to the test group so CI exercises the path. First-call compile is ~50 ms (serial) / ~0.7 s (parallel) cold, ~1 ms / 70 ms warm with numba's on-disk cache — paid once at state reset. fastmath reorders the sum, the same float32-tolerance class as the MLX conv, already covered by the relaxed test tolerances.

New tests cover numba-vs-portable equivalence (spanning the serial/parallel split so both kernels run), float64 dtype preservation, the rail-fill semantics including leading/trailing rails, and the no-numba fallback (via a fixture that hides the kernels).