Johnny-Liou/ReplaySSM

Python

27

14,539 commits

updated Sep 13, 2026

See the code

README

ReplaySSM (a vLLM research fork)

ReplaySSM makes autoregressive and speculative decoding faster for hybrid models (Mamba2, Gated DeltaNet) by caching recent SSM inputs in a small ring buffer instead of writing the recurrent state back to HBM every decode step.

πŸ“ Blog post: ReplaySSM: Cache SSM Inputs, Not State Β Β·Β  πŸš€ Upstreaming to vLLM: RFC #47572, PR #47576

This repository is the original research reference, based on vLLM (Apache-2.0) at upstream commit 37ce34922; the productized version is being contributed upstream via the links above. See the upstream vLLM README below for general install and usage.

Try it out

Scripts under benchmarks/replayssm/ reproduce the blog numbers:

# E2E autoregressive decode speedup: ReplaySSM vs the standard SSM kernel (on single H100)
python benchmarks/replayssm/e2e_decode_speedup.py --model-id nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16 
python benchmarks/replayssm/e2e_decode_speedup.py --model-id Qwen/Qwen3.5-4B --buffer-len 16

# E2E speculative-decode throughput: AR vs standard-spec vs ReplaySSM-spec (on single B300)
python benchmarks/replayssm/e2e_spec_decode_throughput.py --batch-size 512 \
    --model-id nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4
python benchmarks/replayssm/e2e_spec_decode_throughput.py --batch-size 512 \
    --model-id nvidia/Qwen3.5-122B-A10B-NVFP4 --spec-method qwen3_next_mtp

Blackwell / NVFP4 note: the FP4-MoE FlashInfer autotuner can be unstable under CUDA-graph capture on the pre-release Blackwell path, so the benchmark scripts disable it by default (override with --no-disable-flashinfer-autotune).


flagmeaning & constraints
--use-replayssmReplaySSM AR decode kernel. AR only
--use-replayssm-specReplaySSM cached spec kernel.
--replayssm-buffer-len Nring-buffer capacity.
--replayssm-route {output_only,state_and_output}AR compute route (only with --use-replayssm); default output_only.

Core implementation files

ReplaySSM is implemented in Triton kernels.

Kernels β€” Mamba2 (vllm/model_executor/layers/mamba/ops/)

Kernels β€” Gated DeltaNet (GDN) (vllm/model_executor/layers/fla/ops/)

  • fused_recurrent_replayssm.py β€” AR decode. Entry: fused_recurrent_gated_delta_rule_replayssm.
  • gdn_replayssm_spec_decode.py β€” speculative decode on a circular buffer. Entry: gdn_replayssm_spec_decode; cursors commit_gdn_replayssm_spec / reset_gdn_replayssm_spec_cursors.

Everything below is the upstream vLLM README.

vLLM

Easy, fast, and cheap LLM serving for everyone

| Documentation | Blog | Paper | Twitter/X | User Forum | Developer Slack |

πŸ”₯ We have built a vLLM website to help you get started with vLLM. Please visit vllm.ai to learn more. For events, please visit vllm.ai/events to join us.


About

vLLM is a fast and easy-to-use library for LLM inference and serving.

Originally developed in the Sky Computing Lab at UC Berkeley, vLLM has grown into one of the most active open-source AI projects built and maintained by a diverse community of many dozens of academic institutions and companies from over 2000 contributors.

vLLM is fast with:

  • State-of-the-art serving throughput
  • Efficient management of attention key and value memory with PagedAttention
  • Continuous batching of incoming requests, chunked prefill, prefix caching
  • Fast and flexible model execution with piecewise and full CUDA/HIP graphs
  • Quantization: FP8, MXFP8/MXFP4, NVFP4, INT8, INT4, GPTQ/AWQ, GGUF, compressed-tensors, ModelOpt, TorchAO, and more
  • Optimized attention kernels including FlashAttention, FlashInfer, TRTLLM-GEN, FlashMLA, and Triton
  • Optimized GEMM/MoE kernels for various precisions using CUTLASS, TRTLLM-GEN, CuTeDSL
  • Speculative decoding including n-gram, suffix, EAGLE, DFlash
  • Automatic kernel generation and graph-level transformations using torch.compile
  • Disaggregated prefill, decode, and encode

vLLM is flexible and easy to use with:

  • Seamless integration with popular Hugging Face models
  • High-throughput serving with various decoding algorithms, including parallel sampling, beam search, and more
  • Tensor, pipeline, data, expert, and context parallelism for distributed inference
  • Streaming outputs
  • Generation of structured outputs using xgrammar or guidance
  • Tool calling and reasoning parsers
  • OpenAI-compatible API server, plus Anthropic Messages API and gRPC support
  • Efficient multi-LoRA support for dense and MoE layers
  • Support for NVIDIA GPUs, AMD GPUs, and x86/ARM/PowerPC CPUs. Additionally, diverse hardware plugins such as Google TPUs, Intel Gaudi, IBM Spyre, Huawei Ascend, Rebellions NPU, Apple Silicon, MetaX GPU, and more.

vLLM seamlessly supports 200+ model architectures on Hugging Face, including:

  • Decoder-only LLMs (e.g., Llama, Qwen, Gemma)
  • Mixture-of-Expert LLMs (e.g., Mixtral, DeepSeek-V3, Qwen-MoE, GPT-OSS)
  • Hybrid attention and state-space models (e.g., Mamba, Qwen3.5)
  • Multi-modal models (e.g., LLaVA, Qwen-VL, Pixtral)
  • Embedding and retrieval models (e.g., E5-Mistral, GTE, ColBERT)
  • Reward and classification models (e.g., Qwen-Math)

Find the full list of supported models here.

Getting Started

Install vLLM with uv (recommended) or pip:

uv pip install vllm

Or build from source for development.

Visit our documentation to learn more.

Contributing

We welcome and value any contributions and collaborations. Please check out Contributing to vLLM for how to get involved.

Citation

If you use vLLM for your research, please cite our paper:

@inproceedings{kwon2023efficient,
  title={Efficient Memory Management for Large Language Model Serving with PagedAttention},
  author={Woosuk Kwon and Zhuohan Li and Siyuan Zhuang and Ying Sheng and Lianmin Zheng and Cody Hao Yu and Joseph E. Gonzalez and Hao Zhang and Ion Stoica},
  booktitle={Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles},
  year={2023}
}

Contact Us

  • For technical questions and feature requests, please use GitHub Issues
  • For discussing with fellow users, please use the vLLM Forum
  • For coordinating contributions and development, please use Slack
  • For security disclosures, please use GitHub's Security Advisories feature
  • For collaborations and partnerships, please contact us at collaboration@vllm.ai

Media Kit

Contributors

(top 30 of 461)

DarkLight1337

898 commits

WoosukKwon

808 commits

mgoin

545 commits

hmellor

502 commits

Johnny-Liou/ReplaySSM

Python

27

14,539 commits

updated Sep 13, 2026

See the code

README

ReplaySSM (a vLLM research fork)

ReplaySSM makes autoregressive and speculative decoding faster for hybrid models (Mamba2, Gated DeltaNet) by caching recent SSM inputs in a small ring buffer instead of writing the recurrent state back to HBM every decode step.

πŸ“ Blog post: ReplaySSM: Cache SSM Inputs, Not State Β Β·Β  πŸš€ Upstreaming to vLLM: RFC #47572, PR #47576

This repository is the original research reference, based on vLLM (Apache-2.0) at upstream commit 37ce34922; the productized version is being contributed upstream via the links above. See the upstream vLLM README below for general install and usage.

Try it out

Scripts under benchmarks/replayssm/ reproduce the blog numbers:

# E2E autoregressive decode speedup: ReplaySSM vs the standard SSM kernel (on single H100)
python benchmarks/replayssm/e2e_decode_speedup.py --model-id nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16 
python benchmarks/replayssm/e2e_decode_speedup.py --model-id Qwen/Qwen3.5-4B --buffer-len 16

# E2E speculative-decode throughput: AR vs standard-spec vs ReplaySSM-spec (on single B300)
python benchmarks/replayssm/e2e_spec_decode_throughput.py --batch-size 512 \
    --model-id nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4
python benchmarks/replayssm/e2e_spec_decode_throughput.py --batch-size 512 \
    --model-id nvidia/Qwen3.5-122B-A10B-NVFP4 --spec-method qwen3_next_mtp

Blackwell / NVFP4 note: the FP4-MoE FlashInfer autotuner can be unstable under CUDA-graph capture on the pre-release Blackwell path, so the benchmark scripts disable it by default (override with --no-disable-flashinfer-autotune).


flagmeaning & constraints
--use-replayssmReplaySSM AR decode kernel. AR only
--use-replayssm-specReplaySSM cached spec kernel.
--replayssm-buffer-len Nring-buffer capacity.
--replayssm-route {output_only,state_and_output}AR compute route (only with --use-replayssm); default output_only.

Core implementation files

ReplaySSM is implemented in Triton kernels.

Kernels β€” Mamba2 (vllm/model_executor/layers/mamba/ops/)

Kernels β€” Gated DeltaNet (GDN) (vllm/model_executor/layers/fla/ops/)

  • fused_recurrent_replayssm.py β€” AR decode. Entry: fused_recurrent_gated_delta_rule_replayssm.
  • gdn_replayssm_spec_decode.py β€” speculative decode on a circular buffer. Entry: gdn_replayssm_spec_decode; cursors commit_gdn_replayssm_spec / reset_gdn_replayssm_spec_cursors.

Everything below is the upstream vLLM README.

vLLM

Easy, fast, and cheap LLM serving for everyone

| Documentation | Blog | Paper | Twitter/X | User Forum | Developer Slack |

πŸ”₯ We have built a vLLM website to help you get started with vLLM. Please visit vllm.ai to learn more. For events, please visit vllm.ai/events to join us.


About

vLLM is a fast and easy-to-use library for LLM inference and serving.

Originally developed in the Sky Computing Lab at UC Berkeley, vLLM has grown into one of the most active open-source AI projects built and maintained by a diverse community of many dozens of academic institutions and companies from over 2000 contributors.

vLLM is fast with:

  • State-of-the-art serving throughput
  • Efficient management of attention key and value memory with PagedAttention
  • Continuous batching of incoming requests, chunked prefill, prefix caching
  • Fast and flexible model execution with piecewise and full CUDA/HIP graphs
  • Quantization: FP8, MXFP8/MXFP4, NVFP4, INT8, INT4, GPTQ/AWQ, GGUF, compressed-tensors, ModelOpt, TorchAO, and more
  • Optimized attention kernels including FlashAttention, FlashInfer, TRTLLM-GEN, FlashMLA, and Triton
  • Optimized GEMM/MoE kernels for various precisions using CUTLASS, TRTLLM-GEN, CuTeDSL
  • Speculative decoding including n-gram, suffix, EAGLE, DFlash
  • Automatic kernel generation and graph-level transformations using torch.compile
  • Disaggregated prefill, decode, and encode

vLLM is flexible and easy to use with:

  • Seamless integration with popular Hugging Face models
  • High-throughput serving with various decoding algorithms, including parallel sampling, beam search, and more
  • Tensor, pipeline, data, expert, and context parallelism for distributed inference
  • Streaming outputs
  • Generation of structured outputs using xgrammar or guidance
  • Tool calling and reasoning parsers
  • OpenAI-compatible API server, plus Anthropic Messages API and gRPC support
  • Efficient multi-LoRA support for dense and MoE layers
  • Support for NVIDIA GPUs, AMD GPUs, and x86/ARM/PowerPC CPUs. Additionally, diverse hardware plugins such as Google TPUs, Intel Gaudi, IBM Spyre, Huawei Ascend, Rebellions NPU, Apple Silicon, MetaX GPU, and more.

vLLM seamlessly supports 200+ model architectures on Hugging Face, including:

  • Decoder-only LLMs (e.g., Llama, Qwen, Gemma)
  • Mixture-of-Expert LLMs (e.g., Mixtral, DeepSeek-V3, Qwen-MoE, GPT-OSS)
  • Hybrid attention and state-space models (e.g., Mamba, Qwen3.5)
  • Multi-modal models (e.g., LLaVA, Qwen-VL, Pixtral)
  • Embedding and retrieval models (e.g., E5-Mistral, GTE, ColBERT)
  • Reward and classification models (e.g., Qwen-Math)

Find the full list of supported models here.

Getting Started

Install vLLM with uv (recommended) or pip:

uv pip install vllm

Or build from source for development.

Visit our documentation to learn more.

Contributing

We welcome and value any contributions and collaborations. Please check out Contributing to vLLM for how to get involved.

Citation

If you use vLLM for your research, please cite our paper:

@inproceedings{kwon2023efficient,
  title={Efficient Memory Management for Large Language Model Serving with PagedAttention},
  author={Woosuk Kwon and Zhuohan Li and Siyuan Zhuang and Ying Sheng and Lianmin Zheng and Cody Hao Yu and Joseph E. Gonzalez and Hao Zhang and Ion Stoica},
  booktitle={Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles},
  year={2023}
}

Contact Us

  • For technical questions and feature requests, please use GitHub Issues
  • For discussing with fellow users, please use the vLLM Forum
  • For coordinating contributions and development, please use Slack
  • For security disclosures, please use GitHub's Security Advisories feature
  • For collaborations and partnerships, please contact us at collaboration@vllm.ai

Media Kit

Contributors

(top 30 of 461)

DarkLight1337

898 commits

WoosukKwon

808 commits

mgoin

545 commits

hmellor

502 commits

Languages

Python

84.3%

Rust

5.3%

Cuda

5.2%

C++

3.7%