ARM Selective-Scan Kernel for Mamba Inference

The problem

Mamba and other state-space models are supposed to be a good fit for CPU inference — linear time in sequence length, constant memory, none of an attention mechanism’s quadratic cost. In practice, PyTorch’s own implementation falls back to a slow, sequential loop on CPU, with no optimized kernel path available. That gap is what stands between a state-space model and a normal CPU deployment — AWS Graviton, Oracle Ampere, Apple Silicon — that has no GPU to fall back on.

What I built

Working in a 3-person team, I optimized and benchmarked the 1D bidirectional selective scan — one of several scan variants the project covers — implemented as safe Rust with hand-written Arm NEON vectorization and multithreaded execution. The recurrence at the center of a selective scan is sequential by definition, so the speedup comes from getting everything around that recurrence — the elementwise math, the discretization step, the memory layout — onto SIMD lanes and across cores, not from parallelizing the scan itself.

The result: 6.39–8.99× faster than torch.compile on Arm CPUs for that scan.

Verification

Kernel code that silently produces slightly-wrong numbers is worse than kernel code that crashes — a state-space model keeps generating plausible-looking output on a subtly broken scan, and nothing downstream complains. So validation had to check the actual arithmetic at every stage the code passes through, not just that it ran:

• Scalar vs. NEON. The hand-vectorized kernel checked against a plain scalar implementation of the same recurrence before multithreading was introduced as a second variable.
• Single-threaded vs. multithreaded. Confirming that splitting work across cores changed the wall-clock time and nothing else.
• PyTorch integration path. The full path — custom-op registration, tensor marshaling, the actual call into the Rust kernel — checked end to end against high-precision references, not just the math in isolation.

Each layer was validated against a high-precision reference before moving to the next, so a mismatch could be attributed to the layer that introduced it rather than discovered only as a diffuse “the output looks off” at the end.

What I’d do next

• Extend past 1D bidirectional. The team’s kernels now cover additional scan variants and model sizes beyond what I benchmarked directly — the same validate-each-layer approach should carry over cleanly.
• Profile the non-recurrence cost. Once the scan itself is fast, the elementwise steps around it — discretization, gating — start to dominate wall-clock time, which is where the next round of optimization work would go.