Skip to content

v0.6.0: FFT Parallel Scan

Choose a tag to compare

@mrothroc mrothroc released this 24 Apr 04:22
· 626 commits to main since this release

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: RunArch takes TrainOptions struct instead of 13 positional args

Install

brew tap mrothroc/mixlab https://github.com/mrothroc/mixlab
brew install mixlab

Or build from source:

CGO_ENABLED=1 go build -tags mlx -o mixlab ./cmd/mixlab/