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

We thank the reviewer for their time and thoughtful questions. We address the reviewer’s concerns and responses below.

While emphasizing speed-ups, the paper does not quantify the computational cost of offline weight reordering and clustering, which could affect deployment practicality.

We provide the detailed breakdown of the GPU hours on A5000 for our offline clustering, scale calibration, quantization, and weight packing in Table R5.

We note that this is a one-time cost. The quantized model can be deployed to various devices and applications.

The evaluation focuses only on Mamba1/Mamba2. It remains unclear whether Quamba2 generalizes to other SSM architectures (e.g., S4, DSS, DenseMamba).

We apply our method to the latest Jamba-1.6-Mini, a large SSM-based language model with 52 billion parameters (12B active / 52B total), and show the result in Table R12. Following the official guidelines [8], we quantize the Transformer and MoE blocks with bitsandbytes [9, 10], while keeping the Mamba blocks in half-precision. Next, we apply different precision levels to the Mamba blocks using our framework. We observe that the accuracy degradation after quantizing the Mamba blocks to low bit-width is approximately 1%. This demonstrates that our quantization framework enables effective low-bit quantization for large-scale Mamba-Transformer hybrid models, addressing a key limitation of existing LLM quantization techniques.

Offline weight rearrangement may complicate integration into existing inference pipelines, especially for dynamic or frequently updated models.

Our framework only increases the common GPTQ quantization pipeline by 0.05 GPU hours (~3 minutes) on A5000, while largely boosting the accuracy by 13.7% for Quamba2-8B-W4A8, as shown in Table R13.

Notably, our framework also supports LoRA modules for various downstream applications. Once the trained LoRA weights are fused into the main model weights, the same weight reordering and quantization pipeline is reused and applied.

[8] https://huggingface.co/ai21labs/AI21-Jamba-Mini-1.6

[9] Dettmers, et al. "LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale"ArXiv (2022)

[10] Dettmers, et al. "8-bit Optimizers via Block-wise Quantization" ICLR (2022)

We appreciate the reviewer’s positive feedback and respond to the concerns as follows.

Limited evaluation of generalizability, primarily focused on MMLU.

We evaluate our framework on more tasks to show that the W4AX improves generalizability in Table R11. We include BoolQ (accuracy), and the generation-based tasks Natural Questions (NQ) (exact match) and SquadV2 (F1).

We sincerely appreciate the reviewer’spositive feedback and will include the additional evaluations in the final version.

We thank the reviewer for their detailed review and positive feedback. We address the reviewer’s questions below.

a breakdown table for GPU hours.

We provide the detailed breakdown of GPU hours for Quamba2. We report the GPU hours on A5000 for offline clustering, scale calibration, quantization, and weight packing in Table R5.

show the impact of SSD quantization on accuracy.

We show the impact of causal convolution and SSD scan on Mamba2-2.7B in Table R6. We quantize the weights to 4-bit with different activation bit-width settings, and report the accuracy on the Lambada dataset. We apply our quantization techniques to SSD inputs (B, C, x).

The value of quantizing SSD scan and Conv layer is a bit unclear.

The SSD update is a latency and memory bottleneck during generation, especially as batch size increases, although the SSD scan remains minor during prefilling. We profile latency (µs) of the linear layer, causal convolution, and SSD update, along with memory usage (MB) during generation for Quamba2-8B W4A8 and W4A16. As shown in Table R7, cached SSM and convolutional states grow linearly with batch size, eventually exceeding model size and dominating both latency and memory in large-batch scenarios.

Our W4A8 models with 8-bit SSD inputs halve the cached state size and improve time-per-output-token (TPOT, i.e., generation) latency, as shown in Table R8. End-to-end TPOT latency is reported in milliseconds (ms), and OOM indicates an out-of-memory error.

These justify our motivation of quantizing 8-bit SSD and the cached SSM states.

Experiments 5.1 "input SSM" refers to one, some, or all in (A, B, C, x)? Which percentile were used? explaining the motivation or impact would be nice.

For the Mamba2 baseline, we follow the setting in Quamba [5] and use their official implementation. Specifically, we apply the percentile clipping to the x input activation and per-tensor quantization for B and C for Mamba2 models. We list the clipping percentiles in Table R9.

In Table R10, we apply clipping to (B, C, x) for Quamba2-8B-W4A8, respectively, and compare to our proposed weight-reordering and clustering. We apply per-group GPTQ quantization to the weights and Hadamard transforms to the SSD output. We report the accuracy on the Lambada dataset.

We sincerely thank you for your valuable input which prompted us to clarify the details of our work.

[5] Chiang, et al. "Quamba: A Post-Training Quantization Recipe for Selective State Space Models." ICLR (2025).

We thank the reviewer for their insightful comments. We address the reviewer’ concerns and responses below. References [1-3] are from the review.

The paper's novelty is somewhat incremental.

Our SSM quantization framework introduces novel observations, open-sources low bit-width kernels, and explores mixed-precision strategies to improve the generalization, which are distinct from prior SSM quantization work [4, 5]. We summarize the key novelty and contributions of our work:

• We identify three key properties in SSMs: channel order, channel persistence, and state persistence.
• Based on these, we propose sort-and-cluster and per-state-group quantization, validated on large-scale datasets.
• We explore W4AX mixed-precision quantization to boost robustness and generalization of low-bit SSMs.
• Our framework supports W4A8, W4A16, W4AX, and W8A8, enabling flexible deployment and speedups across platforms.

Outperform only two previous methods.

To the best of our knowledge, the most up-to-date, peer-reviewed, and strongest SSM baselines were included by the Jan 30 deadline. We compare against and outperform two and only two latest state-of-the-art SSM quantization methods, MambaQuant and Quamba, which were accepted to ICLR 2025 on Jan 22. We will change “several state-of-the-art” to “two latest state-of-the-art” in our manuscript.

Compared with 1-bit Bi-mamba

We compare our framework against 1-bit Bi-mamba on Mamba2 2.7B in terms of storage, performance, GPU hours, and cost versus the average accuracy of PIQA HS WG ARC-e ARC-c in Table R1. We estimate the A100 GPU hours for a fair comparison.

Although QAT-based Bi-Mamba offers better storage efficiency, it requires several orders of magnitude more tokens, GPU hours, and cost (more than 100,000$\times$) compared to our PTQ framework.

The paper doesn't clearly specify the batch sizes for all accuracy evaluations. Statistical significance testing is absent.

We use a batch size of 16, a fixed random seed, and report the average accuracy over five runs (Section 5.1) for all experiments. We include standard deviation in Table R2 for Quamba2-8B-W4A8. Notably, our method outperforms all baselines on Mamba2 beyond the reported deviations.

The ablation studies focus primarily on the 8B model.

We conduct the same ablation study for Quamba2-2.7B-W4A8, and show the results in Table R3. We note that the group sizes of B/C for 2.7B are one, making it equivalent to per-tensor quantization. Our framework in the last row outperforms all other settings.

Quamba2 better than non-SSM quantized models on edge?

We include the 4-bit Llama-3-8B [7] in Table R4 for reference. Due to time constraints, we report only the Time-To-Last-Token (TTLT) in seconds on an A5000 GPU, using 2K input tokens and 2K generated tokens, along with the average accuracy across six zero-shot tasks. We will include latency results on the Nano 8G in the next version of our manuscript.

Explain the search space and constraints. The 3:1 ratio of W4A8/A16 seems predetermined.

Our search space consists of N^2 configurations, where N denotes the number of layers. We fix the W4A8/A16 ratio and search for the precision for each layer to achieve the best accuracy. We will include more configurations in our manuscript.

Evolutionary search for mixed-precision seen in other domains. Cite relevant on hardware-aware NAS targeting bit-width selection.

Our work to address the performance gap for low bit-width SSMs on large datasets by mixed-precision with generic speedups, which are different from CNNs [2, 3], Transformers [6], and prior SSM quantization work [4, 5]. We will cite and discuss the relevant work in our final version.

[4] Xu, et al. "MambaQuant: Quantizing the Mamba Family with Variance Aligned Rotation Methods." ICLR (2025).

[5] Chiang, et al. "Quamba: A Post-Training Quantization Recipe for Selective State Space Models." ICLR (2025).

[6] Zhao, et al. "Automatic mixed-precision quantization search of bert." IJCAI (2021).

[7] Lin, et al. "QServe: W4A8KV4 Quantization and SystemCo-design for Efficient LLM Serving." MLSYS (2025).