SketchSSM

SketchSSM: Write to the Full State, Read from a Compact Sketch

SketchSSM speeds up decoding in linear-attention and state-space layers (Mamba-2, Gated DeltaNet, KDA). It keeps updating the full recurrent state, but between periodic exact flushes it reads a compact low-rank sketch of the state instead of the full state, which cuts state-read traffic. The sketch size, which sets the traffic reduction, is chosen in the serving configuration; the sketching matrix is computed once per model by offline calibration.

SketchSSM over a window. At a flush step, SketchSSM updates the full state S₀ and refreshes the sketch U. At non-flush steps, it multiplies U by query-dependent coefficients c_t to reconstruct the output; S₀ is not accessed between state updates.

Paper

This repository contains:

• vllm/: vLLM v0.30.0 with SketchSSM.
• sketchssm/kernels/: the SketchSSM CUDA decode kernels.
• sketchssm/calibration/: the offline calibration.
• benchmarks/: per-layer decode 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:

1. --sketchssm: the calibration, a Hub repo id or a local file.
2. --sketchssm-mean-rank: the rank budget, the mean sketch rank per state head (default 10: about 10x less state traffic at standard decode accuracy in the paper).

# Calibration from the Hub
vllm serve nvidia/NVIDIA-Nemotron-Nano-9B-v2 --trust-remote-code \
  --sketchssm ominn/SketchSSM-Nemotron-Nano-9B-v2-BF16 --sketchssm-mean-rank 10 \
  --mamba-ssm-cache-dtype float32 --no-enable-prefix-caching

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

See vllm/README.md for all options.

Benchmarks

See benchmarks/ to measure the per-layer recurrent decode latency of a model 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.