FlashSpeculativeMamba

FlashSpeculativeMamba is a research-oriented implementation of speculative decoding with a fast draft model and a verifier/target model.

Project Layout

FlashSpeculativeMamba/
|-- src/
|   |-- __init__.py
|   |-- speculative_engine.py
|   |-- draft_model.py
|   |-- target_model.py
|   |-- mamba_state_cache.py
|   |-- target_verifier.py
|-- tests/
|   |-- test_engine.py
|-- scripts/
|   |-- check_cuda.py
|-- benchmark.py
|-- demo.py
|-- requirements.txt
|-- .gitignore
|-- README.md

Setup

python -m venv venv
venv\Scripts\activate
pip install -r requirements.txt

Run

python demo.py
python benchmark.py
python -m pytest tests/test_engine.py -v
python scripts/check_cuda.py

Notes

• benchmark.py is the research evaluation loop (formerly main.py).
• demo.py is the project demo entrypoint (formerly riddima/run_demo.py).
• src/mamba_state_cache.py and src/target_verifier.py preserve teammate modules with clearer names.