PaperlayerView exact revision on arXiv

Source-linked AI summary

SSDi8: Accurate and Efficient 8-bit Quantization for State Space Duality

Hyunwoo Kim, Byoungchan Ko, Minseok Kang, Minwoo Kim, Dongjin Lee, Jaehoon Lee, Sungroh Yoon, Dahuin Jung

arXiv:2608.21952v1cs.AI

TL;DR

SSD-based Mamba-2 improves scalability but raises memory and latency overhead, creating a need for architecture-specific compression. SSDi8 introduces persistent INT8 execution with SSD-aware quantization and error correction. It achieves FP16-comparable accuracy with up to 1.4× speedup in W4A8 and W8A8, including validation on Orin NX.

  • Problem

    Mamba-2’s SSD increases memory and latency overhead, while direct application of existing quantization methods can substantially degrade accuracy because SSD has distinctive dimensional and activation distributions.

  • Method

    SSDi8 maintains a persistent INT8 path through activation reuse and sparse-aware reformulation, and reduces quantization error using SSD dimensional decomposition, axis-specific outlier handling, and per-channel mean correction.

  • Results

    SSDi8 achieves accuracy comparable to FP16 and up to 1.4× inference speedup in W4A8 and W8A8, with reported speedups across batch–sequence settings and resource-constrained Orin NX deployment.

  • Takeaways & Limitations

    SSDi8 provides an SSD-specific post-training quantization approach for accurate and efficient inference in large-scale Mamba-2 models and edge environments.

Abstract

from arXiv · show

Recent advances in sequence modeling have highlighted Mamba as a state space architecture offering efficient long-range dependency modeling and providing a viable alternative to Transformers. Building upon this, Mamba-2 introduces the Structured State Space Duality (SSD), which integrates recurrent and attention modes to achieve efficiency and scalability. However, this architectural expansion substantially increases memory and latency overhead, underscoring the need for efficient compression strategies tailored to SSD. In this work, we present SSDi8, the first post-training quantization framework specifically designed for SSD to maintain a persistent INT8 path. SSDi8 introduces a reformulation that decouples element-wise multiplications from matrix multiplications, enabling reuse of quantized activations across modules. Moreover, SSDi8 adaptively quantizes channel-varying activations at cost-effective points, further reducing latency. On the accuracy side, SSDi8 explicitly leverages the intrinsic dimensional decomposition of SSD, exploiting distinct outlier distributions across axes, and incorporates an error correction term based on per-channel error statistics. Comprehensive experiments demonstrate that SSDi8 achieves accuracy comparable to FP16 while delivering up to 1.4x speedup in W4A8 and W8A8 settings. We further validate its robustness in resource-constrained environments by deploying it on the Orin NX device.

1 INTRODUCTION

Mamba-2’s SSD improves scalability but increases memory and latency overhead, while its distinctive computation makes direct Transformer quantization inaccurate. SSDi8 addresses these issues with persistent INT8 execution and SSD-aware accuracy corrections, achieving FP16-comparable accuracy and substantial speedups.

  • Mamba provides efficient long-range dependency modeling, whereas Transformer attention has quadratic computation and memory growth with sequence length.
  • Mamba-2’s SSD integrates recurrent and attention modes, scales beyond 2.7B parameters, and intensifies memory and latency overhead.
  • SSD repeatedly reuses activations and performs sequential updates, causing frequent DRAM accesses and substantial latency overhead.
  • Transformer-oriented quantization methods can substantially degrade SSD accuracy because head and per-head dimensions have different distributions and activations vary by dimension.
  • SSDi8 maintains a persistent INT8 path by reusing quantized activations, quantizing channel-variant and recurrent activations at selected points, and reformulating disruptive element-wise operations.
  • Up to 1.4× inference speedup is achieved in W4A8 and W8A8 while retaining accuracy comparable to FP16; single-inference speedups reach 1.5×.

2 RELATED WORK

Mamba-2 extends state space modeling with structured SSDs, multi-head structure, and scalable state dimensions, while quantization methods remain specialized to different Mamba generations and architectures.

  • Mamba-2’s structured SSD establishes formal equivalence between SSMs and linear attention and enables optimized GEMM-based implementations.
  • Mamba-2 increases the state dimension from Mamba-1’s N = 16 to N = 64 −128 and beyond, while adding a multi-head structure for scalability.
  • These architectural advances enable large-scale parameter expansion but intensify memory and latency overhead, motivating compression and deployment strategies.
  • MambaQuant and Quamba1 target Mamba-1, whereas Quamba2 extends quantization to Mamba-2 using W4A8 and W8A8 settings that include SSD blocks.
  • Quamba2 is limited to SSD-layer inputs and does not adequately address preceding SSD computation.

3 BACKGROUND

Quantization maps continuous tensor values to finite integer levels, while Mamba’s state-space formulation uses hidden-state updates and discretization to process discrete inputs channel by channel.

  • 3.1 QUANTIZATION: Uniform quantization divides a value range into equal intervals and maps each tensor element to its nearest quantized level.
  • 3.1 QUANTIZATION: The quantized tensor uses a scaling factor αX to define the step size from X’s maximum absolute value, together with bit-width b.
  • 3.2 MAMBA-1: Mamba uses a hidden state to efficiently compress and propagate memory in an architecture composed solely of activation operations.
  • 3.2 MAMBA-1: Its state-space equations update the hidden state as h′(t) = Ah(t) + Bx(t) and produce output y(t) = Ch(t).
  • 3.2 MAMBA-1: Zero-Order Hold and a time-step activation ∆ discretize continuous-time matrices A and B for discrete inputs independently along each input channel.
  • Input-dependent activations emphasize important information while suppressing noise, improving long-range dependency modeling.

4 METHODOLOGY

SSDi8 restructures SSD computations to sustain an INT8 execution path, reusing quantized activations while adapting quantization to SSD’s head and channel dimensions. It combines sparse-aware reformulation, per-axis state quantization, recurrent INT8 updates, and mean correction to reduce latency and accuracy loss.

  • SSDi8 workflow: SSDi8 executes much of the SSD pipeline through persistent INT8 representations, reusing quantized activations while retaining selected operations in FP16 when recovery is challenging.The workflow retains dAcs from ChunkCumsum and ChunkScan2 in FP16 because their element-wise multiplications make recovery after quantization difficult.
  • SSD structure: Mamba-2 SSD unifies recurrent and attention modes and represents model activations with independent head and per-head dimensions, where D = H⊙P.SSD partitions sequences into chunks and processes them through ChunkCumsum, ChunkState, StatePassing, ChunkBMM, and ChunkScan modules.
  • Sparse-aware reformulation: SSDi8 moves the element-wise LUTstate multiplication onto X, forming Xscaled before quantization so matrix multiplications can remain on the INT8 path.The reformulation is valid because LUTstate multiplies along the shared l-dimension, while other dimensions operate independently; sparse Xscaled reduces quantization error under the stated conditions.
  • State quantization: SSDi8 quantizes State per-(H,P), allowing common scales across chunks and bit-shift recurrent updates while avoiding quantization along the N axis.StateINT8 can therefore persist through ChunkScan1 for INT8 Tensor Core multiplications, and FP16 Decay is quantized for StatePassing.
  • Efficiency: 1.38× overall speedup and 1.6× ChunkScan speedup are achieved over FP16 by combining persistent recurrent-state INT8, sparse-aware reformulation, and activation reuse.The reported gains arise from the INT8 execution path and reduced intermediate FP16 memory traffic.
  • Error correction: Per-channel mean correction improves accuracy with approximately 1–2% latency overhead by applying corrections only to the output projection layer.A layer-wise sequential update accounts for activation shifts induced by earlier corrections.

5 EXPERIMENTS

Experiments evaluate SSDi8 across zero-shot accuracy, perplexity, latency, resource-constrained deployment, ablations, and a hybrid Mamba–Transformer model. Across these settings, SSDi8 preserves performance relative to FP16 or Quamba2 while reducing latency.

  • Experimental Setup: Experiments compare SSDi8 with FP16, Quamba, and Quamba2 on Mamba-2 models spanning 1.3B, 2.7B, and 8B parameters.Evaluation covers six zero-shot benchmarks and WikiText2 perplexity under W8A8 and W4A8 configurations.
  • Accuracy: 3.9% lower perplexity (7.49 vs. 7.79) under W8A8 and 4.0% lower perplexity (7.62 vs. 7.94) under W4A8 are achieved on the 8B model.SSDi8 consistently narrows the gap to FP16 while outperforming Quamba2 across model scales and bit-widths.
  • Latency: 1.47× speedup over FP16 and 1.38× over Quamba2 are achieved on Mamba-2 2.7B with batch size 32 and sequence length 2048.At the module level, ChunkScan reaches 1.77× over FP16, while StatePassing reaches 2.25×.
  • Resource-Constrained Deployment: SSDi8 consistently outperforms Quamba2 on NVIDIA Orin NX 16G across sequence lengths under W4A8 and W8A8.The evaluation uses Mamba-2 2.7B with batch size 16.
  • Ablation Studies: 67.2% accuracy with SSD quantization rises to 67.4% with mean correction, compared with 51.2% for HadMamba under W4A8 on Lambada.Mean correction adds approximately 1–2% overhead.
  • Hybrid Model: 73.1% FP16 average accuracy becomes 73.0% with INT8-only SSD quantization, while SSD-module latency falls from 19.834 ms to 9.156 ms.End-to-end forward latency decreases from 109.873 ms to 98.904 ms while MLP and attention remain in FP16.

6 CONCLUSION

SSDi8 is an INT8 quantization framework for Mamba-2 SSD that combines persistent INT8 representations, activation reuse, sparse-aware reformulation, and mean correction. Its implementation follows SSD’s chunked state-space computation while targeting accurate and efficient inference.

  • Conclusion: SSDi8 establishes persistent INT8 representations for Mamba-2 SSD through activation reuse and a sparse-aware reformulation.The reformulation addresses element-wise operations that would otherwise disrupt the INT8 execution path.
  • Conclusion: Mean correction compensates for accumulated quantization errors, while internal activation analysis guides quantization strategies for SSD.The framework analyzes SSD’s activation behavior and selects quantization strategies suited to its operations.
  • Conclusion: SSDi8 provides mathematical intuition for sparse-tensor quantization and guidance for settings with element-wise and recurrent operations.The conclusion frames these properties as applicable beyond the specific SSD implementation.
  • Mean Correction: Sequential mean correction updates layer-wise correction terms using statistics from outputs produced by preceding corrected layers.This avoids estimating every correction independently from a fixed initial input distribution.
  • SSD Computation: The SSD layer processes input activations with chunked state-space operations whose final output sums off-diagonal and diagonal contributions.ChunkState and StatePassing handle recurrent information, while ChunkScan1 and ChunkScan2 generate the two output terms.

E ADDITIONAL ACCURACY RESULTS

Additional results show that SSDi8 remains competitive across Mamba-2 scales and quantization methods, while approaching FP16 perplexity on the Pile benchmark. These evaluations extend the reported accuracy evidence beyond the primary tables.

  • Additional Accuracy Results: SSDi8 matches Quamba2 under W4A16 quantization even when using W4A8 quantization.Table 8 evaluates 1.3B, 2.7B, and 8B Mamba-2 models across six zero-shot tasks.
  • Pile Perplexity: SSDi8 surpasses Quamba2 and approaches FP16-level performance under W8A8 on Pile perplexity for 1.3B and 2.7B models.The evaluation uses sequence length L = 2048.

F ADDITIONAL ABLATION STUDIES

Ablations examine activation quantization axes and mean-correction placement, identifying SSD-specific accuracy sensitivities and a latency–accuracy trade-off. The selected design reflects these observed differences.

  • Quantization Axis: Per-G,N quantization performs best for activations B and C, with only a 0.02 difference from per-G quantization.X and State show substantial degradation when either the P or H axis is omitted.
  • Mean Correction Placement: Mean correction produces its highest accuracy gain immediately after SSD layers, but applying it at the out-projection layer yields the lowest latency.The final placement balances accuracy improvement against latency.

G ADDITIONAL LATENCY AND MODEL SIZE RESULTS

SSDi8 maintains memory usage nearly identical to Quamba2 while reducing SSD-module latency relative to FP16 and Quamba2 at evaluated batch sizes.

  • Memory usage is nearly identical to Quamba2, with SSDi8 requiring 2.953GB versus 2.948GB for the 2.7B model under W8A8.For the 8B model, usage is 9.867GB versus 9.860GB, a +0.07% increase.
  • At batch 16, SSDi8 reduces SSD-module latency from 14.745ms with FP16 to 10.646ms, a 27.8% reduction.
  • At batch 32, SSDi8 reaches 19.000ms versus 25.797ms for Quamba2, a 26.3% latency reduction, while FP16 encounters OOM.OOM denotes Out-Of-Memory.

H LONGER CONTEXT RESULTS

In long-context evaluation from 2k to 14k tokens, SSDi8 improves throughput and SSD latency while maintaining stable perplexity where FP16 and Quamba2 become unstable.

  • Across 2k–14k tokens, SSDi8 consistently achieves higher throughput and lower SSD latency than FP16 and Quamba2, with the gap widening as context length increases.
  • At 14k tokens, SSDi8 improves throughput from 11,362 tokens/s with FP16 to 14,200 tokens/s, a gain of +2,838 tokens/s (+25.0%).
  • At 14k tokens, SSDi8 reduces SSD latency from 30.870 ms to 21.138 ms, a 9.73 ms reduction (-31.5%).
  • SSDi8 maintains PPL within 8.94–9.36 across 2k–14k tokens, while FP16 rises from 13.845 at 4k to 113.407 at 6k and Quamba2 shows similar instability.

I CALIBRATION SENSITIVITY ANALYSIS

SSDi8 is relatively insensitive to calibration sample count and dataset choice, while module-wise analysis identifies SSD as a major optimization target after projection layers.

  • Calibration robustness: Average accuracy across six zero-shot tasks remains largely stable as calibration samples increase from 128 to 2048.Individual benchmarks show only minor differences.
  • Calibration robustness: Changing the calibration dataset produces comparable performance, with no single calibration corpus consistently resulting in lower accuracy.
  • Latency breakdown: Projection layers are the dominant FP16 latency bottleneck, followed by SSD, while normalization and other components contribute relatively little.SSDi8 applies Hadamard-based activation quantization and GPTQ weight quantization to projections while focusing its primary optimization on SSD.

K BATCH-SIZE SENSITIVITY ANALYSIS

Across batch sizes, SSDi8 reduces SSD latency and improves throughput, with larger batches amplifying the benefits of INT8 execution.

  • Batch-size sensitivity: SSDi8 consistently reduces SSD-module latency and improves throughput across batch sizes spanning edge-oriented small batches to large-batch serving.
  • Batch-size sensitivity: At batch size 16, SSD latency decreases from 2.757 ms to 2.671 ms, a 3.1% reduction, while throughput improves by 3.8%.
  • Batch-size sensitivity: At batch size 128, SSD latency is reduced by 25.5% and throughput increases by 19.5%.
  • Batch-size sensitivity: At batch size 256, SSDi8 achieves a 32.6% reduction in SSD latency, from 35.194 ms to 23.730 ms, and a 21.8% throughput improvement.
  • Interpretation: The benefits of quantizing ChunkState and ChunkBMM become increasingly pronounced under higher computational intensity.Larger batches amplify the arithmetic and memory-efficiency advantages of INT8 execution across edge and cloud-serving regimes.
Loading 2608.21952v1…