Quamba2: A Robust and Scalable Post-training Quantization Framework for Selective State Space Models (2503.22879v3)

Abstract: State Space Models (SSMs) are emerging as a compelling alternative to Transformers because of their consistent memory usage and high performance. Despite this, scaling up SSMs on cloud services or limited-resource devices is challenging due to their storage requirements and computational power. To overcome this, quantizing SSMs with low bit-width data formats can reduce model size and benefit from hardware acceleration. As SSMs are prone to quantization-induced errors, recent efforts have focused on optimizing a particular model or bit-width for efficiency without sacrificing performance. However, distinct bit-width configurations are essential for different scenarios, like W4A8 for boosting large-batch decoding speed, and W4A16 for enhancing generation speed in short prompt applications for a single user. To this end, we present Quamba2, compatible with W8A8, W4A8, and W4A16 for both Mamba1 and Mamba2 backbones, addressing the growing demand for SSM deployment on various platforms. Based on the channel order preserving and activation persistence of SSMs, we propose an offline approach to quantize inputs of a linear recurrence in 8-bit by sorting and clustering for input $x$, combined with a per-state-group quantization for input-dependent parameters $B$ and $C$. To ensure compute-invariance in the SSM output, we rearrange weights offline according to the clustering sequence. The experiments show that Quamba2-8B outperforms two state-of-the-art SSM quantization methods and delivers 1.3$\times$ and 3$\times$ speed-ups in the pre-filling and generation stages, respectively, while offering 4$\times$ memory reduction with only a $1.6\%$ average accuracy drop. The evaluation on MMLU shows the generalizability and robustness of our framework. The code and quantized models will be released at: https://github.com/enyac-group/Quamba.

Overview of "Quamba2: A Robust and Scalable Post-training Quantization Framework for Selective State Space Models"

The paper introduces Quamba2, a post-training quantization framework specifically designed for Selective State Space Models (SSMs). These models have emerged as a promising alternative to Transformers due to their memory efficiency and high performance in tasks such as LLMing, vision, and audio processing. However, deploying large SSMs on platforms with limited resources presents significant challenges related to storage requirements and computational demands. Quamba2 addresses these challenges by providing a robust framework that efficiently reduces model size through quantization, thus enabling hardware acceleration.

Key Contributions and Methodology

1. Quantization Configurations: Quamba2 supports multiple bit-width configurations (W8A8, W4A8, and W4A16) for both Mamba1 and Mamba2 backbones, catering to varied deployment needs. The framework utilizes channel order preservation, activation persistence, and offline weight and input reordering techniques.
2. Sort-and-cluster approach: To minimize quantization-induced errors in SSM inputs, Quamba2 leverages the order-preserving and persistent activation properties of SSMs. The input channels are sorted by their calibrated maxima and then clustered, allowing smoother value ranges within groups and improving quantization precision.
3. Per-state-group quantization: For input-dependent parameters (B, C) that exhibit activation persistence, Quamba2 applies quantization per state group. This technique enhances precision by using group-specific scaling factors, effectively mitigating performance loss due to lower bit-width quantization.
4. Efficient Kernel Design: The framework includes optimized CUDA kernels to support 4/8-bit operations in embedding layers, SSM blocks, and output layers, resulting in significant latency and memory reductions across diverse platforms.

Results and Implications

Quamba2 demonstrates substantial improvements over existing SSM quantization methods. For instance, it achieves up to 3x speed-up in generation tasks alongside a 4x reduction in memory usage with merely a 1.6% drop in accuracy across various zero-shot tasks. Additionally, Quamba2's compatibility with different bit-width configurations provides flexibility for deploying SSMs in cloud environments and edge devices. The evaluation on the MMLU dataset underscores the generalizability and robustness of the framework, even with heavy quantization settings.

Future Directions

The paper suggests several avenues for future research. One potential direction is exploring advanced quantization techniques for further improving the precision and efficacy of low-bit models in complex reasoning tasks. Another area of interest could be extending Quamba2 to support more SSM variants or hybrid models that integrate SSM features with other neural architectures. Additionally, investigating adaptive quantization techniques that dynamically adjust precision based on computational or memory constraints could enhance the framework's applicability to a broader range of scenarios.

In conclusion, Quamba2 represents a significant advancement in the field of efficient model deployment, particularly for SSMs. It addresses critical challenges related to resource limitations while maintaining high performance and adaptability across varied application domains. The framework not only paves the way for wider adoption of SSMs but also contributes to the broader discourse on scaling machine learning models in constrained environments.