PerfMamba: Performance Analysis and Pruning of Selective State Space Models
This work addresses performance bottlenecks for deploying selective SSMs in real-world applications, offering incremental improvements through empirical analysis and pruning.
The paper tackled the lack of understanding in runtime behavior and resource utilization of selective state space models (SSMs) like Mamba-1 and Mamba-2, finding that the SSM component consumes significant computational resources and proposing a pruning technique that achieves a 1.14x speedup and reduces memory usage by 11.50% while maintaining accuracy.
Recent advances in sequence modeling have introduced selective SSMs as promising alternatives to Transformer architectures, offering theoretical computational efficiency and sequence processing advantages. A comprehensive understanding of selective SSMs in runtime behavior, resource utilization patterns, and scaling characteristics still remains unexplored, thus obstructing their optimal deployment and further architectural improvements. This paper presents a thorough empirical study of Mamba-1 and Mamba-2, systematically profiled for performance to assess the design principles that contribute to their efficiency in state-space modeling. A detailed analysis of computation patterns, memory access, I/O characteristics, and scaling properties was performed for sequence lengths ranging from 64 to 16384 tokens. Our findings show that the SSM component, a central part of the selective SSM architecture, demands a significant portion of computational resources compared to other components in the Mamba block. Based on these insights, we propose a pruning technique that selectively removes low-activity states within the SSM component, achieving measurable throughput and memory gains while maintaining accuracy within a moderate pruning regime. This approach results in performance improvements across varying sequence lengths, achieving a 1.14x speedup and reducing memory usage by 11.50\%. These results offer valuable guidance for designing more efficient SSM architectures that can be applied to a wide range of real-world applications.