Mamba: Linear-Time Sequence Modeling with Selective State Spaces

3.3 Efficient Implementation of Selective SSMs

1. How does Mamba compute a time-varying recurrence efficiently on GPUs?
   With kernel fusion (load parameters from slow HBM into fast SRAM, discretize and run the recurrence there, write only outputs back), a work-efficient parallel scan instead of a sequential loop, and recomputation of intermediate states in the backward pass instead of storing them; memory use matches FlashAttention Transformers.
   • Gu & Dao, 2024, Mamba (First Conference on Language Modeling (COLM 2024)), §3.3