sketchssm 0.1.15

| 🌐 Project Page | 📄 Paper | 🤗 Calibration Dataset | 📦 PyPI |

Overview

Hybrid-attention models replace most softmax attention layers with linear attention (Mamba-2, Gated DeltaNet, KDA), reducing KV-cache growth and enabling larger decode batches, where recurrent-state access becomes a major bottleneck. Reducing the state size by quantization or pruning cuts this traffic, but its approximation errors propagate and accumulate through subsequent decode steps. To reduce state-update traffic, ReplaySSM buffers keys and values over a window of W steps and applies their accumulated updates to the full state once per window. However, each new query still requires a full-state read, even though the state remains unchanged between state updates.

SketchSSM keeps the full-state updates and approximates the reads:

• Flush step (every W steps): read the full state S₀ once, apply the buffered updates from the ring buffer, compute the compact sketch U from the updated state, and write back the state and the sketch.
• Non-flush steps: combine the sketch U and the coefficient map C with the query to reconstruct the output, without reading S₀.

The sketching matrix is computed once per model by offline calibration. The sketch size (the mean sketch rank per state head) is chosen as a serving configuration and sets the trade-off between traffic reduction and accuracy. Across Mamba-2, GDN and KDA models, SketchSSM at mean sketch rank 8 reduces state-access traffic by about 10x while matching the average accuracy of the full-state baseline.

Demo

Nemotron Nano 9B v2 on one RTX PRO 6000 Blackwell decodes IFEval, MATH-500, HumanEval and MBPP prompts (lm-eval, 4 samples each at T=0.6) as one stream at batch 320, without and with ReplaySSM or SketchSSM (mean rank 8, W=16). SketchSSM reaches 7,634 output tokens/s: 2.26x the baseline and 1.64x ReplaySSM, with the same average accuracy (76.2 ± 1.0 vs 76.6 and 76.2, each ± 0.9). Prefill is computed ahead of time; the video shows decode at 8x speed.

This repository contains:

• vllm/: vLLM v0.30.0 with SketchSSM.
• sketchssm/kernels/: the SketchSSM CUDA decode kernels.
• sketchssm/calibration/: the offline calibration.
• evaluation/: accuracy, decode throughput and per-layer linear-attention latency in vLLM.

Installation

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

Quick start: serve with vLLM

SketchSSM needs a calibration file (calibration.pt) made by offline calibration. Calibrations for the models below are on the Hugging Face Hub. To calibrate a new model yourself, see sketchssm/calibration/.

_{See all calibrations.}

Enable SketchSSM with --sketchssm and a calibration, a Hub repo id or a local file:

# Calibration from the Hub
vllm serve nvidia/NVIDIA-Nemotron-Nano-9B-v2 --trust-remote-code \
  --sketchssm SketchSSM/Nemotron-Nano-9B-v2-BF16 \
  --mamba-ssm-cache-dtype float32

# Local calibration file
vllm serve nvidia/NVIDIA-Nemotron-Nano-9B-v2 --trust-remote-code \
  --sketchssm outputs/nano/calibration.pt \
  --mamba-ssm-cache-dtype float32

Optionally, --sketchssm-mean-rank sets the rank budget, the mean sketch rank per state head (default 8: about 10x less state traffic at accuracy comparable to the full-state baseline).

See vllm/README.md for all options.

Mamba prefix caching works in align mode on Mamba-2 models (the default for Nemotron when prefix caching is on); for Gated DeltaNet and KDA models add --no-enable-prefix-caching.

vLLM support

SketchSSM is being upstreamed to vLLM. Until a feature is merged, use the fork in vllm/.

Evaluation

Accuracy versus state access traffic (paper, Figure 5). Accuracy on MATH-500, AIME25, GPQA Diamond and LiveCodeBench for Nemotron Nano v2, Nemotron 3 Super, Qwen3.8 Flash-Next and GLM 5.3 Flash, against the reduction in state access traffic relative to Standard. At mean rank 8 (W=16, 9.4-10.7x less state traffic), SketchSSM matches Standard's average accuracy on all four models, while pruning (GHOST, DRRQR) and quantization (DSQ) lose accuracy at much smaller reductions.

Linear-attention latency on one NVIDIA B300 (paper, Figure 7; W=16, batch 128, 256 and 512): (a) a non-flush step, (b) a flush step, (c) the speedup over Standard across a whole window. At mean rank 8 and batch 512, SketchSSM is 7.30x faster than Standard on Nemotron 3 Super (Mamba-2), 5.02x on Qwen3.8 Flash-Next (GDN) and 5.24x on GLM 5.3 Flash (KDA), and 3.24x, 3.43x and 3.69x faster than ReplaySSM.

See evaluation/ to measure a model's accuracy (IFEval, MATH-500, HumanEval, MBPP), its decode throughput at a chosen batch size, and the per-layer recurrent (linear-attention) decode latency in vLLM, and sketchssm/kernels/ for kernel tuning and microbenchmarks.

Citation

If you use SketchSSM in your research, please cite:

@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},
}

License

SketchSSM is released under the Apache License 2.0. The vllm/ directory is a fork of vLLM, also under Apache-2.0. Calibration files are derived from the base models' weights and are also subject to those models' licenses.