Research note 4
Recurrent-state quantization for linear attention: a first GPU kernel
Abstract
Hybrid models replace part of their attention with linear attention, which keeps a fixed-size recurrent state for each head in place of a growing KV cache. For Qwen3.5-9B that state is 50 MB per request in fp32, and it is read and rewritten at every token. LeapQuant [1] stores it in 8 bits by rounding once per window of tokens and keeping the largest errors in a few high-precision compensators. We built a reference implementation and a first GPU kernel. The reference reproduces the method’s arithmetic. The kernel is 7.2 times slower than SGLang’s window kernel at batch 512, because the fold that closes each window takes 73 ms against a budget near 1 ms. We report the profile and the redesign.
Background
A standard attention layer keeps every past token and consults all of them for each new one. A linear-attention layer keeps instead one matrix per head that summarizes everything seen so far [2]. That matrix is the recurrent state. In Qwen3.5-9B, 24 of the 32 layers are of this kind, each with 32 heads, and each head holds a matrix of 128 by 128 numbers. One request therefore owns 24 × 32 × 16,384 numbers, which is 50.3 MB in fp32.
For every generated token each head decays the matrix, reads what it holds for the current key, forms a correction, writes the correction back as a rank-one update, and reads the output. The arithmetic is small. The cost is moving the matrix, which is loaded from GPU memory and stored back at every token. The kernel is bound by memory bandwidth and not by compute (technical note 2).
Storing the state in fewer bytes attacks that cost directly, and it also raises the number of requests and cached prefixes that fit in memory.
Why rounding the state is difficult
A KV cache entry is written once and only read thereafter, so it can be rounded once [4]. A recurrent state is rewritten at every token. If it is rounded after every token, each token starts from an already rounded matrix and adds a new error, and the errors compound over thousands of tokens. The LeapQuant paper measures the effect: with a bf16 state rounded at every token, the AIME score of Qwen3.5-9B falls from 87.9 to 72.1 [1].
LeapQuant makes two changes. It rounds less often: the low-precision state is held fixed for a window of 16 tokens, each token’s change is kept as a small record in higher precision, and outputs are computed from the fixed state together with the records. At the end of the window a fold applies the records and rounds once. And it rounds less harmfully: before rounding, the largest outliers of the state are set aside as a few rank-one compensators kept in higher precision, and the remainder is smoothed. A related method reaches lower bit widths by a different route [3].
What we built
- A reference implementation in C++, written to be read. It follows the equations of the paper and defines the numbers a GPU kernel must reproduce.
- A recorder and a replay tool. Nine traces of 4,095 tokens each were recorded from three layers under three kinds of prompt, and can be replayed through the reference and through a kernel.
- A GPU kernel with two parts: an append step, which adds one record and computes the output, and a fold, which rebuilds the state, selects compensators, rescales, rounds and empties the buffer.
Results
Numerical behavior
All nine traces pass the reference’s self-checks. On six of them the 8-bit state gives lower output error than an ideal bf16 state held over the same window. On the remaining three it gives higher error, and all three are from layer 30, where the worst head reaches 15.4% output error. In layer 0 the four compensators reduce the error about a hundredfold. In layer 30 they change it from 3.63% to 3.13%.
One trace has been verified on the GPU. Across its 510 folds the kernel’s state matched the reference in every one of 267,386,880 entries.
Speed
| Batch | SGLang, packed kernel | SGLang, window kernel | Our kernel | Our kernel, with gate arithmetic | Ratio to window kernel |
|---|---|---|---|---|---|
| 32 | 0.0821 | 0.0583 | 0.3523 | 0.5981 | 10.3 |
| 128 | 0.3070 | 0.2010 | 1.4137 | 1.5653 | 7.8 |
| 512 | 1.2063 | 0.7869 | 5.3483 | 5.6800 | 7.2 |
The ratio compares our kernel with gate arithmetic to SGLang’s window kernel.
Our kernel is slower than SGLang’s window kernel at every batch size recorded.
Where the time goes
A step is defined as fifteen appends and one fold, divided by sixteen. At batch 512 the profiler attributes 73.4 ms to one fold. Spread over sixteen tokens that is 4.59 ms of a 5.44 ms step, or 84%. The fitting of the compensators takes 57.6 ms of the fold, 78.5%.
The same definition gives the fold a budget. At batch 512 an append takes 0.714 ms and the window kernel averages 0.787 ms, so to equal the window kernel a fold may cost 1.17 ms more than an append. The recorded figure is 74.2 ms, about 64 times the budget. At batch 128 the append alone is already slower than the window kernel.
Fewer sweeps would not be enough. If the compensator fit vanished entirely, the step would still take about 1.84 ms, more than twice the baseline.
| Cause | Mechanism |
|---|---|
| Bank conflicts in shared memory | Each lane of a warp walks one row of the matrix, so the 32 lanes read addresses 128 words apart. These fall in one memory bank and are served one after another. A pass costs about 1.21 ms in place of 0.145 ms, and a fold makes 44 such passes: about 47 of the 73 ms. |
| Number of passes | Four compensators are found one at a time, with eight sweeps each: 72 walks over the matrix. |
| Loads inside the inner loop | The rebuild reads from global GPU memory twenty times for each matrix entry. |
| Serial sums | Each sum is taken by one thread in a fixed order while the others wait. |
| Barriers | A fold passes about 211 points at which all threads wait for each other. |
| Occupancy | A block needs about 65 KB of shared memory, so at most three blocks run on one streaming multiprocessor. |
| Division | The rounding divides twice for each entry. |
The bank-conflict estimate is derived from the code and the profile. It awaits confirmation by a new run.
The fit itself has a weakness independent of speed. It starts from the longest row of the matrix on the assumption that this row overlaps the strongest pattern. A counterexample shows otherwise: scored against the strongest pattern, that start captures 16 of 291, where a block fit from a random start captures all 291.
Redesign
The following changes are planned. None is complete.
- Lay the matrix out in shared memory with a stride of 129 words, which by our arithmetic removes the bank conflict.
- Fit the compensators together as a block, starting from the previous window’s compensators.
- Determine why layer 30 loses, and set the policy for each layer separately.
- Make the compressed state the only state of an active request. At present the fp32 pool is kept beside it, which adds about 46% to state memory where it should save.
- Measure serving speed and task accuracy together. Neither has been measured yet.
The first target is a mean layer time at least 20% below the fastest baseline that passes the quality checks, with gate arithmetic and complete folds included. On the figures above that is 0.63 ms or less at batch 512, against 0.787 ms.
What the paper reports
The paper measures B200, RTX PRO 6000 and RTX 5090 GPUs, and reports a kernel speedup of 2.05 to 3.70 times over an fp32 kernel, a decode-step speedup of 1.22 to 1.37 times for whole models at batch 512, an end-to-end speedup of 1.47 times, and 41 to 56% less memory with prefix caching [1]. It has no measurement on the H100 generation. We therefore do not take its kernel gains as an acceptance figure for our hardware. Our target is set against the fastest baseline that passes the quality checks on our workload.
Prefix caching for models of this kind raises its own questions, since a recurrent state cannot be truncated to an earlier position as a KV cache can [5].
Limitations
- All timings use an fp32 state. No timing with a bf16 state exists yet.
- One of nine traces has been verified on the GPU, and it is the easiest layer.
- Replay error is an arithmetic measure. It is not task accuracy.
- No end-to-end serving result exists.
References
- Y. Pan, H. Xi, K. Zhu, et al. LeapQuant: Efficient Linear Attention with Accurate Recurrent State Quantization. 2026.
- S. Yang, J. Kautz, A. Hatamizadeh. Gated Delta Networks: Improving Mamba2 with Delta Rule. 2024.
- B. Yao, H. Xu, H. Lin, et al. STEPQuant: When and Where Errors Matter in Delta-Rule Recurrent State Quantization. 2026.
- Z. Liu, J. Yuan, H. Jin, et al. KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache. 2024.
- R. Pan, Z. Wang, Z. Jia, et al. Marconi: Prefix Caching for the Era of Hybrid LLMs. 2024.
Write to us about this note.