SketchSSM

Write to the Full State,
Read from a Compact Sketch

Omin Kwon1 JoongWon Shin1Minseo Kim2Kurt Keutzer2 Sehoon Kim3◆ Jae W. Lee1◆

1Seoul National University2UC Berkeley3KAIST

Co-corresponding authors

Seoul National University UC Berkeley KAIST

Same model. Less memory traffic. Faster decoding.

SketchSSM in two minutes: the decode demo, the bottleneck, the key idea, and the results.
Demo: Nemotron Nano 9B v2 · 1× RTX PRO 6000 Blackwell · Batch 320, sped up; prefill is computed ahead of time. Watch the full decode demo at 8× playback.
7,634 tokens/s

SketchSSM decode throughput

2.26×

faster than the standard baseline

1.64×

faster than ReplaySSM

Demo results
Method Accuracy (%) Throughput
(tokens/s)		 Decode speedup
IFEval MATH-500 HumanEval MBPP Average vs Standard vs ReplaySSM
Standard 86.1 ± 1.3 67.7 ± 1.7 81.4 ± 2.5 71.1 ± 1.9 76.6 ± 0.9 3,383 1.00× --
ReplaySSM 85.8 ± 1.3 66.4 ± 1.7 81.9 ± 2.4 70.8 ± 1.9 76.2 ± 0.9 4,669 1.38× 1.00×
SketchSSM 85.4 ± 1.3 67.5 ± 1.7 80.9 ± 2.6 70.9 ± 1.9 76.2 ± 1.0 7,634 2.26× 1.64×

Accuracy is reported as mean ± standard error over problems. Speedups and throughput are from the demo.

Benchmarks
IFEval · MATH-500 · HumanEval · MBPP
Sampling
4 samples / prompt · Temperature 0.6 · lm-eval
SketchSSM
Mean sketch rank 8 · Window 16 steps

01 / The idea

Keep the state full.
Make the reads small.

A recurrent state carries information into every future step. Compressing that state can carry approximation errors forward, too. SketchSSM preserves full-state updates and approximates only the readout to avoid propagating state-compression errors into subsequent decoding steps.*

How state-compression errors carry forward

Each update inherits the previous state's error. Compressing the updated state adds a new error term.

Full stateSt = MtSt−1 + Bt

Compressed stateS̃t = Q(MtS̃t−1 + Bt)

Error propagationEt = MtEt−1 + εt

Et = S̃t − St. Mt is the state transition, Bt the new update, Q state compression, and εt the new compression error. This comparison holds the input sequence and update terms fixed to isolate direct state-compression error.

An early compression error survives subsequent state updates From left to right, four states are read to produce output vectors and updated along the decoding timeline. Starting with zero state error, the error terms become epsilon one, then M two times epsilon one plus epsilon two, then M three times M two times epsilon one plus M three times epsilon two plus epsilon three. Colored bands distinguish propagated error terms from newly introduced compression errors. Output read S₀ No error t = 0 update + compress+ ε₁ Output read S̃₁ ε₁ t = 1 update + compress+ ε₂ Output read S̃₂ M₂ε₁ + ε₂ t = 2 update + compress+ ε₃ Output read S̃₃ M₃M₂ε₁ + M₃ε₂ + ε₃ t = 3 Decode time →

State-write access: Already amortized across multiple decoding steps by ReplaySSM, making full-state updates affordable while avoiding state-compression error propagation.

State-read access: Reduced by reading from a compact sketch refreshed alongside each full-state update.

Once per window · Flush

Write to the full state

Apply buffered updates to the uncompressed state, then refresh a compact sketch. The state is never replaced with its approximation.

Avoid propagating state-compression errors.*

Between flushes · Non-flush

Read from a small sketch

Combine sketch vectors with query-dependent coefficients to approximate the state readout. Recent buffered updates contribute separately to the output.

No full-state access between flush steps.

* Readout approximation errors can still propagate through output vectors and affect subsequent full-state updates. Here, avoiding error propagation refers specifically to preventing the accumulation of errors from information discarded by state compression during ring-buffered state updates.

HBM and SM accesses over a SketchSSM window. Flush steps read and update the full state and refresh the sketch and coefficient map. Non-flush steps combine buffered key-value updates with the sketch readout without accessing the full state.
One full-state update every W steps; compact-sketch reads in between. The sketching basis is calibrated offline, while the sketch and coefficient map are refreshed at each flush.

02 / How to Make a Sketch

Preserve the output.
Build a compact query basis.

Within each window, the full state stays fixed, but the effective query changes at every decoding step. Computing each query's output directly therefore requires rereading the same full state.

Our idea is to approximate these changing queries with a small, fixed query basis. A single full-state read at the flush step precomputes the outputs for all basis vectors, forming a compact sketch. Subsequent queries combine these precomputed outputs instead of rereading the full state.

Different queries. One fixed basis.

Queries at different decode steps represented in a fixed three-basis subspace Three colored axes are the offline-calibrated basis vectors omega one, omega two, and omega three. Gray query vectors at different decode steps use different coefficients in this same fixed basis. The three-dimensional view is schematic, not an ordinary Euclidean projection. ω₁ ω₂ ω₃ q̃₁q̃₂q̃₃ New query at each decode step
q~t≈q~^t=ct,1ω1+ct,2ω2+ct,3ω3

Precompute the basis outputs.

One full-state read produces a much narrower sketch The full state transpose is a wide V by K matrix. Multiply it by the fixed K by G basis to produce the compact V by G sketch. State and sketch share the same output dimension V, while the sketch contains far fewer columns. Each sketch column u i equals S zero transpose times omega i. At each flush · One full-state read Same full state S₀ᵀ V × K × Fixed basis Ω ω₁ ω₂ ω₃ K × G Compact sketch U u₁ u₂ u₃ V × G uᵢ = S₀ᵀωᵢ

U = [u₁, u₂, u₃] = S0⊤Ω

Each decode step: change the coefficients, reuse the stored outputs. No full-state reads needed.
otstate≈o^tstate=S0⊤q~^t=S0⊤(ct,1ω1+ct,2ω2+ct,3ω3) =ct,1u1+ct,2u2+ct,3u3=Uct

To choose a basis that preserves the output, we minimize output reconstruction error—equivalently, state-weighted query approximation error.

ℰt=∥S0⊤q~t−Uct∥22=∥H12(q~t−Ωct)∥22

H = S0S0⊤; S0 is the full state at the window boundary, and q̃t is the effective query incorporating within-window transitions.

Offline · Fixed

Sketching matrix

ℒG(Ω)=𝔼𝒟[minc∥E012(q~t−Ωc)∥22]
Ω⋆=E0−12PG

Use the average state geometry E0 = 𝔼𝒟[S0S0⊤]. Uncentered PCA of zt = E01/2q̃t gives PG, the top G eigenvectors of 𝔼𝒟[ztzt⊤]. Calibrate once per state head; keep the basis fixed during inference.

Per window · Flush

Sketch & coefficient map

U=S0⊤Ω⋆
C=(U⊤U)†U⊤S0⊤

Refresh U and the state-dependent map C with each full-state update. Sketch construction is fused into the update kernel. Here, † denotes the pseudoinverse; the implementation uses an efficient approximation of C.

Per step · Non-flush

Coefficient vector & readout

ct=Cq~t,o^tstate=Uct

Choose query-dependent coefficients that minimize ∥S0⊤q̃t − Uc∥22, then combine the sketch vectors. Read only the compact sketch and map; the buffered contribution to the output remains exact.

See the paper for derivations and implementation details ↗

03 / Accuracy

About 10× less state traffic.
Accuracy largely preserved.

At mean sketch rank Ḡ = 8, SketchSSM preserves average accuracy across four benchmarks for all four models spanning Mamba-2, Gated DeltaNet, and KDA, while reducing state-access traffic by 9.4–10.7× relative to Standard.

Follow the green curves. Moving right means less state traffic; staying high means better accuracy. Pruning and quantization lose accuracy at smaller traffic reductions because errors introduced by state compression propagate across decoding steps.

Accuracy and generation length versus state-access traffic reduction for four models on MATH-500, AIME25, GPQA Diamond, and LiveCodeBench. SketchSSM uses mean ranks 8, 4, and 2; rank 8 preserves average accuracy at 9.4–10.7 times less state-access traffic.
Paper, Figure 5. Accuracy and verbosity versus state-access traffic reduction relative to Standard as detailed in the paper’s appendix, “State Memory Traffic Reduction.” ReplaySSM and SketchSSM use W = 16. DSQ uses {8, 6, 4}-bit states; GHOST (Nano/Super/GLM) and DRRQR (Qwen) prune {50, 62.5, 75}% of the state. SketchSSM uses Ḡ ∈ {8, 4, 2} for all four models. At Ḡ = 8, which reduces state-access traffic by 9.4–10.7×, SketchSSM preserves average accuracy across the four benchmarks for all four models, whereas pruning and quantization degrade accuracy at smaller reductions. Verbosity is the mean generation length normalized to Standard. Raw accuracy results are provided in the paper’s appendix, “Accuracy Results on Decode Benchmarks.” View full-size figure ↗

Sketch rank controls the accuracy–traffic trade-off. Larger traffic reductions can decrease accuracy.

04 / Linear-attention kernel speedup

Smaller state-access,
Faster Linear Attention

Most steps read only the sketch, not the full state, to reduce memory traffic.

At a flush step, SketchSSM refreshes the sketch by fusing the sketch construction computation into the existing state update kernel.

7.30×

Mamba-2

Nemotron 3 Super

5.02×

Gated DeltaNet

Qwen3.8 Flash-Next

5.24×

KDA

GLM 5.3 Flash

Whole-window kernel speedups over Standard · Mean sketch rank Ḡ = 8 · 1× NVIDIA B300 · W = 16 · Batch 512

Linear-attention performance on NVIDIA B300 at mean sketch ranks 8 and 4. Columns show non-flush latency, flush latency, and whole-window speedup over Standard at batch sizes 128, 256, and 512. At rank 8 and batch 512, speedups are 7.30, 5.02, and 5.24 times for Mamba-2, GDN, and KDA.
Paper, Figure 6. Linear-attention latency on one NVIDIA B300 with W = 16 and batch sizes 128, 256, and 512. Rows show Nemotron 3 Super, Qwen3.8 Flash-Next, and GLM 5.3 Flash, with SketchSSM at Ḡ ∈ {8, 4}. (a) Non-flush-step latency, incurred at each of the W − 1 steps between flush steps. (b) Flush-step latency, incurred once per window. Hatching indicates sketch and coefficient-map construction overhead at the flush step. (c) Total linear-attention speedup over Standard across the W steps of a window. Lower latency is better in (a–b); higher speedup is better in (c). View full-size figure ↗

In the paper’s end-to-end evaluation at mean sketch rank Ḡ = 8, these kernel improvements translate to 2.77× higher decode throughput at 2K context and 2.33× at 8K on Nemotron 3 Super with one NVIDIA B300, compared with Standard at each method’s maximum batch size.

Try SketchSSM

Bring SketchSSM
to your serving stack

vLLM integration, CUDA kernels, evaluation scripts, and ready-to-use calibration files.

Get started on GitHub 

git clone https://github.com/SNU-ARC/SketchSSM.git
cd SketchSSM/vllm
VLLM_USE_PRECOMPILED=1 python -m pip install -e .
python -m pip install sketchssm   # CUDA kernels; without it vLLM uses its Triton kernels

Calibrated query bases are available.

Stay tuned for more models. You can also create your own offline calibration for your model using the SketchSSM code repository.

SketchSSMGLM 5.3 FlashNVFP4 SketchSSMQwen3.8 Flash-NextNVFP4 SketchSSMQwen3.5 9BBF16 SketchSSMNemotron 3 SuperNVFP4 SketchSSMNemotron Nano 9B v2BF16

Citation

@misc{kwon2026sketchssmwritestateread,
  title={SketchSSM: Write to the Full State, Read from a Compact Sketch},
  author={Omin Kwon and JoongWon Shin and Minseo Kim and Kurt Keutzer and Sehoon Kim and Jae W. Lee},
  year={2026},
  eprint={2609.33051},
  archivePrefix={arXiv},
  primaryClass={cs.LG},
  url={https://arxiv.org/abs/2609.33051}
}