Repository navigation
v0.6.0: FFT Parallel Scan
FFT-Based Parallel Scan
Replaces the sequential T-step loop in OP_SCAN with an FFT-based causal
convolution. The gate is constant across time (shape [D]), so the
recurrence h[t] = gate*h[t-1] + (1-gate)*x[t] is equivalent to
convolving with exponential kernel k[τ] = (1-gate) * gate^τ.
40x throughput improvement for mamba and RWKV blocks at long sequences:
- Mamba 3L T=1024: 77K tok/s (was ~2K with sequential scan)
- Falls back to sequential for T ≤ 32
Other Changes
- fix: tok/s and ETA reporting uses wall-clock average (fixes pipelined training measurement)
- refactor:
RunArchtakesTrainOptionsstruct instead of 13 positional args
Install
brew tap mrothroc/mixlab https://github.com/mrothroc/mixlab
brew install mixlabOr build from source:
CGO_ENABLED=1 go build -tags mlx -o mixlab ./cmd/mixlab/