ARM-Scan — Optimized Mamba Selective-Scan Kernel
An Arm-optimized selective-scan kernel for Mamba, fusing the forward and reverse recurrences and tuned for NEON SIMD and multicore execution on Graviton4.
- ~1.7×
- Scan fusion speedup
- 18.9×
- SISO speedup
- 38.7×
- Multicore scaling
The problem
Mamba’s selective-scan operator is the part of the architecture that doesn’t parallelize like attention does — it’s an inherently sequential recurrence, which makes it the natural bottleneck once everything else is fast. Running it bidirectionally (forward and reverse) usually means computing the recurrence twice. The question was whether the two directions actually need two independent computations, or whether some of that work is redundant.
Approach
The kernel is written in Rust and hand-tuned for Arm: NEON SIMD for data-parallel work within a core, and multicore execution to split the sequence-length dimension across cores, targeting AWS Graviton4. The core architectural idea is fusing the 1D bidirectional scan so the forward and reverse recurrences share discretization — the step that turns Mamba’s continuous-time parameters into the discrete-time form the recurrence actually runs — instead of recomputing it independently for each direction.
What broke
There wasn’t a single dramatic failure here so much as a correctness constraint that shaped the whole design: sharing discretization across two directions only pays off if the fused result is numerically identical to running both scans independently. Every version of the fused kernel was checked against a reference PyTorch implementation before being accepted, which meant the ~1.7× speedup from fusion had to be proven free, not just fast.
Results
| Metric | Result |
|---|---|
| Shared discretization | ~1.7× speedup over independent forward/reverse scans |
| Single-core, single-in-single-out | Up to 18.9× faster than baseline |
| Multicore | Up to 38.7× scaling on Graviton4 |
What I’d do differently
Pending — this project’s retrospective wasn’t in the source material this case study was drafted from. Krithik: what would you change if you did this again?