STEPQuant: When and Where Errors Matter in Delta-Rule Recurrent State Quantization

Summary (Overview)

• Problem: Recurrent states in Delta-rule linear attention models (e.g., Qwen3.8-27B, Kimi-Linear-48B) become a memory bottleneck during concurrent serving, yet direct low-bit quantization causes severe accuracy degradation due to error propagation through successive state updates.

• Key Insight: Quantization error impact depends on two complementary dimensions: temporal (errors in long-lived memory persist across many decoding steps) and spatial (errors in different key rows affect model outputs differently, with state magnitudes varying along both rows and columns).

• Proposed Method: STEPQuant combines Lifetime-aware Bit Allocation (temporal) with Key-Row-Aware Dual-Axis Fitting (spatial) to allocate precision based on error magnitude, memory lifetime, and key-row impact on output error.

• Results: Under a nominal 6-bit budget, STEPQuant closely matches FP32-state accuracy on both models; even at 4 bits, it outperforms uniform INT8. Integrated into SGLang, 6-bit STEPQuant achieves 5.03× recurrent-state compression and reduces total serving memory by up to 68.7%.

• Implementation: Code available at https://github.com/Dreamer-Toby/STEPQuant with optimized GPU kernels integrated into SGLang.

---

Introduction and Theoretical Foundation

Background

Unlike conventional softmax attention (which maintains a growing KV cache), linear attention summarizes past tokens into a fixed-size recurrent state matrix. Hybrid models like Qwen3.8-27B and Kimi-Linear-48B combine gated Delta-rule recurrent memory with standard attention.

Memory bottleneck: Although state size is fixed per request, each concurrent request requires a separate persistent state. In official SGLang deployment, the FP32 state pool of Qwen exceeds the memory footprint of its BF16 weights at 70 concurrent requests.

State Update Formulation

For a single head, the state after token $t$ is $S_t \in \mathbb{R}^{d_k \times d_v}$. The gated Delta-rule update is:

$$S_t = D_t S_{t-1} + \beta_t k_t \left(v_t^\top - k_t^\top D_t S_{t-1}\right) = (I - \beta_t k_t k_t^\top) D_t S_{t-1} + \beta_t k_t v_t^\top, \quad y_t = S_t^\top q_t \tag{1}$$

where $D_t$ controls memory retention, $\beta_t \in [0,1]$ controls write strength, and $y_t$ is the head output (readout).

The two architectures differ in their retention gate:

• GDN (Qwen): scalar gate per head, $D_t = \alpha_t I$
• KDA (Kimi): channel-wise gates, $D_t = \text{diag}(d_{t,1}, \ldots, d_{t,d_k})$

Both can be unified as:

$$S_t = A_t S_{t-1} + B_t, \quad \text{where} \quad A_t = (I - \beta_t k_t k_t^\top) D_t, \quad B_t = \beta_t k_t v_t^\top \tag{2}$$

Quantization Error Propagation

With quantized states, the update becomes:

$$X_t = A_t \hat{S}_{t-1} + B_t, \quad \hat{y}_t = X_t^\top q_t, \quad \hat{S}_t = \mathcal{Q}_t(X_t) \tag{4}$$

where $\mathcal{Q}_t$ quantizes $X_t$ and returns its dequantized approximation $\hat{S}_t$. Symmetric uniform quantization is defined as:

$$\mathcal{Q}_{b,s}(x) = s \cdot \text{clip}\left(\text{round}\left(\frac{x}{s}\right), -q_b, q_b\right), \qquad q_b = 2^{b-1} - 1 \tag{3}$$

---

Methodology

Temporal Dimension: Lifetime-Aware Bit Allocation

Error Propagation Analysis

Proposition 1 (Conditional error propagation): For identical inputs and gates, let $E_t = \widehat{S}_t - S_t$ denote accumulated error and $\varepsilon_t = \mathcal{Q}_t(X_t) - X_t$ the quantization error added at step $t$. Then:

$$E_t = (I - \beta_t k_t k_t^\top) D_t E_{t-1} + \varepsilon_t = A_t E_{t-1} + \varepsilon_t, \qquad \widehat{y}_t - y_t = E_{t-1}^\top A_t^\top q_t \tag{5}$$

If $\|k_t\|_2 \leq 1$, $0 \leq \beta_t \leq 1$, and $0 \preceq D_t \preceq I$, then $\|A_t\|_2 \leq \|D_t\|_2 \leq 1$.

Key findings:

• Previously accumulated error propagates through $A_t$; the retention gate $D_t$ attenuates it, and the Delta update reduces its component along $k_t$ while leaving orthogonal components unchanged
• Errors in directions rarely aligned with subsequent keys depend mainly on gate decay
• When retention is close to one, errors can persist for many decoding steps
• Empirical validation: Heads with longer gate half-lives exhibit larger accumulated state error (Spearman's $\rho_S \approx 0.80$)

Bit Allocation Objective

For each unit $u$ (entire head in Qwen, key row in KDA), estimate:

• $d_u(b)$: reconstruction distortion at $b$ bits
• $\ell_u$: mean log retention over calibration tokens

The error retention factor after $j$ updates is approximated by:

$$\prod_{s=1}^{j} r_{t+s,u} \approx \exp(j \ell_u) \tag{6}$$

The lifetime weight over $H$ steps:

$$L_u = \sum_{j=0}^{H-1} \exp(2j \ell_u) \tag{7}$$

Given average bit budget $\bar{b}$, select bit widths to minimize lifetime-weighted distortion:

$$\min_{b_u \in \mathcal{B}_{\bar{b}}} \sum_u L_u d_u(b_u) \quad \text{s.t.} \quad \sum_u n_u b_u \leq \bar{b} \sum_u n_u \tag{8}$$

Highest-risk units are retained in FP16 as sparse pivots.

Spatial Dimension: Key-Row-Aware Dual-Axis Fitting

Key-Row Impact on Readout Error

From Eq. (5), the readout error is:

$$\Delta y_t = \hat{y}_t - y_t = E_{t-1}^\top A_t^\top q_t = \sum_i (A_t^\top q_t)_i E_{t-1,i,:}^\top \tag{9}$$

Define $g_t = A_t^\top q_t \in \mathbb{R}^{d_k}$, where $g_{t,i}$ weights the contribution of error in key row $i$. The row-impact score is:

$$\omega_i = \mathbb{E}_{\text{cal}}[g_{t,i}^2] \tag{10}$$

Empirical validation: Groups with larger $\omega$ produce greater perplexity degradation when quantized to INT4.

Two-Axis State Geometry

• Recurrent states exhibit large-magnitude outliers along both key rows and value columns
• Maximum-to-median RMS contrasts: 10.3× (key rows) and 19.4× (value columns)
• Outliers persist throughout decoding (98.6% of sampled states exceed 3× contrast)

Dual-Axis Fitting

Represent the updated state $X = X_t$ as:

$$\hat{X}_{ij} = r_i c_j z_{ij} \tag{11}$$

where $r_i > 0$ and $c_j > 0$ are scale factors for key row $i$ and value column $j$, and $z_{ij}$ is the low-bit integer.

Row factors (accounting for both magnitude and impact):

$$r_i = m_i^{1/2} w_i^{-1/2} \tag{13}$$

where $m_i = \frac{1}{d_v} \sum_j |X_{ij}|$ is the row magnitude.

Column factors (minimizing impact-weighted reconstruction error):

$$\min_{\{c_j > 0\}} \sum_{i,j} w_i^2 \left(X_{ij} - r_i c_j z_{ij}\right)^2 \tag{14}$$

Kernel Implementation in SGLang

• Offline: Lifetime-Aware Bit Allocation and FP16 pivot selection (no per-token overhead)
• Fused kernel: tilewise state reconstruction + Delta update + readout in a single pass
• Overlapped execution: Key-Row-Aware Dual-Axis Fitting and packed writeback run on a separate CUDA stream

---

Empirical Validation / Results

Experimental Setup

• Models: Qwen3.8-27B (GDN), Kimi-Linear-48B-A3B-Instruct (KDA)
• Hardware: Four NVIDIA A800 GPUs
• Baselines: FP32 state, uniform INT4/6/8
• Calibration: 32 WikiText-2 segments (2048 tokens each)
• Benchmarks: 7 long-generation reasoning tasks, 6 short-generation understanding tasks

Long-Generation Results (Table 1, BF16 weights)

Short-Generation Results (Table 2, BF16 weights)

Compatibility with W4 Weights (Table 3)

With 4-bit AWQ weights, 6-bit STEPQuant achieves 79.27% (Qwen) and 58.62% (Kimi), only 0.05 and 0.33 points below FP32-state baselines.

Component Ablation (Table 4, Qwen, 4-bit budget)

Key ablation findings:

• Spatial fitting: Outperforms Q-Mamba's DSQ by 66.31 points at 4 bits (73.95% vs. 7.64%)
• Pivot protection: Protecting only 1.39% of Qwen heads in FP16 improves 4-bit average by 6.70 points
• Combined effect: STEPQuant outperforms both individual components, demonstrating complementarity

Serving Efficiency

At batch size 512 on Qwen with W4 weights:

• Total memory: Reduced from 419.73 to 131.18 GiB (68.7% reduction)
• Recurrent-state memory: 80.1% reduction (5.03× compression)
• State-update time: 65.6% reduction (2.91× faster)

Generation Length Analysis

Uniform quantization substantially increases output length (e.g., Kimi generates 63.40K tokens on AIME at 4-bit, approaching the 65536-token limit with near-zero accuracy). STEPQuant maintains output lengths close to FP32-state levels, consistent with preserved accuracy.

---

Theoretical and Practical Implications

Theoretical Contributions

1. Formal error propagation analysis: Proposition 1 provides a rigorous characterization of how quantization errors propagate through gated Delta-rule updates, showing that retention gates and Delta updates jointly shape error persistence.

2. Temporal-spatial decomposition: The paper establishes that quantization error impact depends on both when (memory lifetime) and where (key-row position) errors occur, providing a principled framework for state quantization.

3. Lifetime-weighted optimization: The bit allocation objective (Eq. 8) formalizes how to trade off quantization distortion against error persistence, generalizing prior single-axis approaches.

Practical Implications

1. Memory-efficient serving: STEPQuant enables sub-8-bit recurrent state quantization with negligible accuracy loss, directly addressing the concurrency-driven memory bottleneck in hybrid linear attention models.

2. Deployment-ready integration: The SGLang integration with optimized kernels demonstrates practical feasibility, including overlapped computation and packed-state storage.

3. Complementarity with weight quantization: STEPQuant remains effective with 4-bit AWQ weights, making it suitable for fully quantized deployment scenarios.

4. Comparison to prior work: STEPQuant achieves 6-bit quantization with negligible degradation, while concurrent work DAMP reports preserved accuracy at 9.9 bits per state value—demonstrating that lower precision is feasible with careful spatial-temporal design.

---

Conclusion

STEPQuant addresses the challenge of low-bit quantization for Delta-rule recurrent states by recognizing and exploiting two complementary error dimensions:

• Temporally, quantization errors persist according to memory lifetime, motivating Lifetime-aware Bit Allocation that assigns higher precision to units with larger and longer-lived errors
• Spatially, key rows differ in readout impact and state magnitudes vary along both axes, motivating Key-Row-Aware Dual-Axis Fitting with impact-weighted scales

The combined approach achieves near-FP32 accuracy at 6 bits and outperforms uniform INT8 at 4 bits on both Qwen3.8-27B and Kimi-Linear-48B across long- and short-generation benchmarks. With SGLang integration, STEPQuant delivers substantial memory savings (up to 68.7% total memory reduction) and faster state updates (2.91×), making it a practical solution for memory-efficient concurrent serving of hybrid linear attention models.

Future directions suggested by this work include extending the spatial-temporal quantization framework to other state-space architectures, exploring even lower bit budgets with adaptive pivoting strategies, and investigating the interaction between state quantization and other compression techniques (e.g., activation quantization, speculative decoding).