Diving into the Implementation Details of Mamba

We continue our in-depth examination of SSM-based models. At the beginning of the seminar, we will review the Mamba architecture—what “selectivity” means and how state updates are implemented—then go through minimal implementations of S4 and Mamba. After that, we will dive into the details of efficient implementation according to the following plan.

• Brief introduction to GPU computing: GPU architecture and the model of massive parallelism.
• The prefix sum (Scan) problem: formulation and the evolution of GPU algorithms for solving it, from Sequential Scan to Work-Efficient Parallel Scan.
• We will analyze where and how Parallel Scan is applied in Mamba, as well as other GPU optimizations used during training.