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

AI summary

Overview

Research area: Efficient LLM inference — post-training quantization of recurrent states in Delta-rule linear attention models (natural language processing / systems).

Technical level: Advanced. The paper combines a formal error-propagation analysis of gated Delta-rule state updates with a mixed-precision quantization method and custom GPU kernels integrated into a serving engine.

Scope: The paper analyzes how quantization errors in fixed-size recurrent states accumulate over decoding time and distribute across state dimensions, then proposes a spatial-temporal quantization framework to compress those states during serving.

What This Paper Is About

Linear attention models replace the growing KV cache with a fixed-size recurrent state, but each concurrent request needs its own persistent state, so total state memory grows with batch size even for short contexts. Naively quantizing those states to low precision breaks accuracy badly, because each Delta-rule update feeds the previous quantization error back into the next state. The paper's goal is to quantize recurrent states below 8 bits while keeping accuracy close to the FP32-state baseline, and to make that compression fast enough for real serving.

Key Contributions

1. Error analysis across two dimensions. The authors analyze how quantization errors propagate through gated Delta-rule updates both temporally (long-lived memory units accumulate larger errors) and spatially (errors in different key rows affect the readout differently, and state magnitudes vary along both key rows and value columns).

2. Lifetime-aware Bit Allocation. A mixed-precision scheme that assigns bit widths to state units based on both their reconstruction distortion at a given bit width and their mean log retention (memory lifetime), under a fixed average bit budget, with a small number of highest-risk units retained in FP16 as sparse pivots.

3. Key-Row-Aware Dual-Axis Fitting. A spatial method that assigns separate scale factors to key rows and value columns, incorporating a calibrated row-impact score so that rows whose errors matter more to the output receive finer resolution.

4. Serving implementation. Packed-state kernels integrated into SGLang that fuse state reconstruction, the Delta update, and readout, with scale fitting and packed writeback overlapped on a separate CUDA stream. Reported at over 5x recurrent-state compression and up to 2.91x faster state updates.

Main Findings

• Temporal error accumulation is lifetime-dependent. Because the transition matrix A_t satisfies a spectral norm bound of at most 1 under stated conditions, accumulated error persists mainly in directions rarely aligned with subsequent keys. Across 2304 Qwen heads under uniform INT6 state quantization on C4, heads with longer gate half-lives show larger accumulated state error, with Spearman's rho approximately 0.80.

• Uniform quantization degrades sharply. On Qwen3.8-27B long-generation benchmarks, mean accuracy falls from 80.60 under FP32 states to 71.86 (INT8), 45.04 (INT6), and 12.73 (INT4). On Kimi-Linear-48B-A3B-Instruct it falls from 61.52 to 56.02 (INT8), 45.70 (INT6), and 21.63 (INT4).

• 6-bit STEPQuant matches FP32 states. STEPQuant@6bit reaches 80.59 on Qwen and 61.47 on Kimi across the seven long-generation tasks, essentially matching the FP32-state baselines of 80.60 and 61.52.

• 4-bit STEPQuant beats uniform INT8. STEPQuant@4bit averages 80.51 on Qwen and 58.52 on Kimi, above the uniform INT8 results (71.86 and 56.02) and close to FP32 states.

• Short-generation results hold up. On six language understanding tasks, 4-bit STEPQuant trails FP32 by only 0.15 points on Qwen (87.63 vs. 87.78) and 0.25 points on Kimi (68.11 vs. 68.36), whereas uniform INT4 drops to 65.74 and 42.24 respectively.

• Spatial structure is two-axis. Maximum-to-median RMS contrasts along key rows and value columns are 10.3x and 19.4x; both exceed 3x in 98.6% of sampled states, and the outliers persist across 2048 decoding steps.

• Both components are complementary. On three long-reasoning tasks at a nominal 4-bit budget, STEPQuant achieves 84.72 versus 73.95 for spatial only, 12.87 for temporal only, 7.64 for Q-Mamba's dual-axis scaling adaptation, and 3.97 for uniform INT4, against an FP32 average of 84.61.

• FP16 pivots matter. Protecting only 1.39% of Qwen heads with FP16 pivots improves the 4-bit three-task average by 6.70 points and 6-bit AIME accuracy by 14.16 points.

• Aggressive uniform quantization produces overthinking. Under 4-bit quantization, Kimi generates an average of 63.40K tokens on AIME and 64.22K on HMMT, approaching the 65536-token limit while achieving near-zero accuracy on both tasks.

• Serving savings. With W4 weights at batch size 512, STEPQuant@6bit reduces Qwen total serving memory from 419.73 to 131.18 GiB (68.7%) and reduces recurrent-state memory by 80.1% (5.03x compression) and state-update time by 65.6% (2.91x faster). For Kimi, the reported compression is 5.08x with a 53.7% total memory reduction.

• Compatibility with quantized weights. With 4-bit AWQ weights, 6-bit STEPQuant averages 79.27 on Qwen and 58.62 on Kimi, only 0.05 and 0.33 points below the corresponding FP32-state baselines.

Methodology in Plain English

The authors start from the observation that a Delta-rule state is updated and then re-quantized at every decoding step, so errors do not just appear once — they get carried forward and added to. They write down the update rule in a form that separates the part that transports old memory from the part that writes new memory, then show that the error at any step equals the transported previous error plus the new quantization error.

From that analysis they conclude the problem has two sides. Temporally, an error written into a memory unit only fades as fast as that unit's retention gate lets it fade, so units that keep information longer deserve more bits. Spatially, an error sitting in one key row can matter much more to the output than an equally sized error in another row, and the state matrix has large outliers along both key rows and value columns at once, so scaling along a single axis is not enough.

Their method therefore does two things. First, during offline calibration on WikiText-2 segments, it measures each state unit's quantization distortion at candidate bit widths and its mean log retention, then solves a small optimization that spends a fixed average bit budget where errors are both large and long-lived; a tiny set of the riskiest units is kept in FP16. Second, it gives every key row a scale that grows with the row's average magnitude but shrinks with a row-impact score (the expected squared weight of that row's error in the readout), and fits value-column scales to minimize the impact-weighted reconstruction error. At inference, a fused kernel reconstructs state tiles, applies the Delta update, and produces the readout in one pass, with the scale refitting and packed writeback overlapped with the rest of the model's computation.

Why This Matters

Recurrent states are a fixed-size but concurrency-scaling memory cost, and the paper notes that in official SGLang deployment the FP32 state pool of Qwen exceeds the memory footprint of its BF16 weights at 70 supported concurrent requests. Showing that these states can go to 6 bits — or even 4 bits in these evaluations — without losing accuracy changes what reallocation is possible in a serving cluster.

Real-world applications:

• High-concurrency LLM serving, where batch size is limited by per-request state memory rather than by weights.
• Long-context deployments of hybrid linear-attention models, where each active session holds a persistent state.
• On-device or edge inference of hybrid attention models, where a few bits per state value can determine whether a model fits.
• Reasoning workloads with very long generations (up to 65536 tokens per sample here), where state memory is held for a long time.

Industry relevance: The method is implemented as packed-state kernels inside SGLang with measured memory and latency numbers on A800 GPUs, and it composes with 4-bit AWQ weight quantization. That combination targets the practical regime where weight memory has already been squeezed and state memory becomes the dominant serving overhead.

Future Directions

• Pushing below 4 bits. The paper shows 4-bit states are feasible on these models; whether the lifetime and row-impact machinery extends to 3-bit or lower, and with what calibration cost, is not reported.
• Generalizing beyond GDN and KDA. The unit of allocation is defined as an entire head in Qwen and a key row in KDA. How the scheme should be defined for other recurrent architectures — including Mamba-style state-space models, where the authors adapt Q-Mamba's dual-axis scaling as a baseline — remains open.
• Calibration efficiency and robustness. Bit allocation, pivots, and row-impact scores are fixed offline from 32 WikiText-2 segments of 2048 tokens each. Whether these statistics transfer across domains, languages, or extremely long contexts is not reported.
• Comparison with concurrent work at matched budgets. The paper compares against DAMP, which reports preserved accuracy at 9.9 bits per state value in its evaluated settings, but a direct head-to-head at equal precision and equal serving configuration is left to future work.

Target Audience

Researchers and engineers working on LLM inference efficiency, particularly those dealing with linear attention, hybrid attention, or state-space models. It is most useful to readers who already understand quantization basics (scales, bit widths, post-training quantization) and who care about the serving-systems side — kernel design, memory accounting, and batching — rather than only model accuracy. Practitioners deploying Qwen3.8-27B or Kimi-Linear-48B-A3B-Instruct on SGLang-style runtimes are the most direct beneficiaries.

Authors’ abstract

Linear attention replaces growing KV caches with fixed-size recurrent states, yet these persistent states can become a substantial memory bottleneck under concurrent serving. Directly quantizing recurrent states to low precision often leads to severe accuracy degradation, as quantization errors propagate through successive state updates. We discover that the impact of these errors depends on two complementary dimensions: temporally, errors in long-lived memory can persist across many decoding steps; spatially, errors in different key rows affect model outputs differently, while state magnitudes vary substantially along both rows and columns. Motivated by these observations, we propose STEPQuant, a spatial-temporal post-training quantization framework for Delta-rule recurrent states. STEPQuant allocates precision according to error magnitude and memory lifetime, and jointly fits key-row and value-column scales based on state distributions and key-row impact on output error. Experiments on Qwen3.8-27B and Kimi-Linear-48B-A3B-Instruct across both long- and short-generation benchmarks show that STEPQuant closely matches FP32-state accuracy under a nominal 6-bit budget and outperforms uniform INT8 in its 4-bit configuration. Integrated into SGLang with optimized GPU kernels, 6-bit STEPQuant achieves over 5x recurrent-state compression and reduces total serving memory by up to 68.7%. Our code is available at https://github.com/Dreamer-Toby/STEPQuant.