Optimize MLX kernel reuse and dense rate reductions

[cboulay]

Summary

• Stop specializing the threshold-crossing Metal kernels on chunk dimensions and refractory width.
• Read dimensions from MLX-generated shape metadata and pass remaining per-call values as runtime inputs, so jittered chunk lengths reuse one compilation.
• Replace the refractory bitmap's manual least-significant-bit search with Metal ctz and bits &= bits - 1.
• Use direct reshape/reduction for uniform completed rate bins, retaining cumulative sums for irregular fractional boundaries and single carried-bin completions.

Why

The previous Metal templates included N_SAMPLES, N_CHANNELS, N_WORDS, and REFRAC_WIDTH. Every unseen chunk shape therefore produced a new kernel specialization and a large latency spike. The dense rate path also materialized a full cumulative-sum tensor and gather indices even when its completed bins were regular.

Measurements

M4 Pro, MLX 0.31.2, explicit evaluation and synchronization:

• Before: each unseen chunk length incurred roughly 44–78 ms of compilation latency.
• After: one initial compilation remains; subsequent unseen lengths took 0.23–0.68 ms.
• Uniform completed-bin rate processing improved by approximately 17–36%, depending on shape.
• Single-bin completions with prior carry retain the previous cumulative-sum path after profiling showed that case should not use the direct reducer.

The ctz scan was neutral for sparse crossings and substantially faster for crossing-heavy inputs in the exploratory benchmark.

Validation

• uv run pytest -q: 155 passed
• uv run ruff check .: passed
• uv run ruff format --check .: passed
• git diff --check: passed