Skip to content
Ravi-VKPublic

About

No description, website, or topics provided.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Latest commit

 

History

66 Commits

Folders and files

Repository files navigation

SSR: Self-Speculation for Faster Reasoning Models

SSR overview

This repository contains the code used to reproduce the latency experiments for Self-Speculation for Faster Reasoning Models.

SSR is a training-free speculative decoding method for reasoning language models. Instead of using a separate draft model, SSR asks the same model to draft an answer from a partial chain-of-thought (CoT) while the original request continues reasoning. Once the full CoT is available, the draft is verified under the higher-budget context. Accepted draft tokens reduce the number of sequential decode steps needed for the final answer.

The implementation in this repository builds on a patched vLLM tree and combines two reuse mechanisms:

  • Prefix verification: accept the longest draft prefix supported by the full-CoT verifier.
  • Suffix decoding: keep using the draft as an external suffix source during continuation, so useful spans after the first mismatch are not discarded.

The code is intended as a research artifact for reproducing the reported latency results. The accompanying paper will be released separately. This is not a production serving release.

Results At A Glance

The main experiments evaluate SSR on long-form and structured generation tasks with 4B-class reasoning models. The largest measured gain is on ClassEval with Gemma-4-E4B-it, where SSR reduces end-to-end latency by about 24%.

Model Benchmark n Latency reduction
Qwen3.5-4B LongProc 2K 989 9.1%
Qwen3.5-4B ClassEval 100 18.5%
Qwen3.5-4B HumanEval 164 2.9%
Gemma-4-E4B-it LongProc 2K 989 7.1%
Gemma-4-E4B-it ClassEval 100 24.1%
Gemma-4-E4B-it HumanEval 164 14.6%

The Gemma ClassEval ablation shows why suffix reuse matters:

Prefix verification Suffix decoding Speedup
yes no 1.086x
no yes 1.318x
yes yes 1.318x

How SSR Works

Reasoning models usually produce a CoT before the final response. SSR uses this structure as a speculation opportunity:

  1. The parent request generates CoT tokens normally.
  2. At a configured draft budget, the scheduler forks a child request from the same prefix.
  3. The child appends the model-specific end-of-thinking marker and drafts answer tokens from the partial CoT.
  4. The parent continues reasoning to the full budget.
  5. When reasoning finishes, the final-CoT verifier accepts a prefix of the draft and then continues generation.
  6. During continuation, vLLM suffix decoding can use the remaining draft tokens as an external suffix cache source.

This adds compute during the overlap window, but much of the draft latency is hidden behind ongoing CoT generation. Prefix caching keeps the forked child request from recomputing the shared prompt and partial-CoT prefix.

Method Components

Partial-CoT drafting

SSR uses one model at two reasoning budgets. A lower-budget partial CoT drafts an answer; the completed CoT verifies it. This avoids auxiliary draft models, draft heads, or target-model-specific training.

Prefix verification

The verifier checks the draft answer under the full-CoT context and accepts the longest valid prefix. This is especially useful for structured outputs such as code, where imports, class definitions, function signatures, and broad structure often stabilize before every detail is resolved.

Draft-backed suffix decoding

Long outputs can diverge early and then share long spans later. Prefix-only verification would discard those later spans. SSR therefore feeds the draft into vLLM V1 suffix decoding as an external source during continuation.

Iterative drafting

The best draft point is not known in advance because the final CoT length is not known in advance. The experimental code supports spawning drafts at multiple CoT intervals and using earlier drafts to help produce later drafts.

Repository Layout

specdec/
|-- README.md                         # this guide
|-- environment-specdecsd.yml         # conda environment snapshot
|-- requirements-specdecsd.txt        # pip package snapshot
|-- data/
|   |-- README.md                     # dataset download notes
|   `-- REPLICATION_DATA.md           # exact data artifacts used by scripts
|-- images/
|   `-- self-spec-compact.png         # method overview figure
|-- notes/
|   |-- experiment_plan.md            # experiment checklist
|   `-- experiment_results.md         # local replication summaries
|-- scripts/
|   |-- prepare_replication_data.sh   # one-command data preparation
|   |-- verify_replication_data.py    # sanity checks for prepared data
|   |-- prepare_code_eval_sources.py  # HumanEval/ClassEval source builder
|   |-- benchmarks/                   # latency benchmark runners
|   |-- analysis/                     # result summarizers
|   `-- inference/                    # vLLM JSONL inference helpers
`-- vllm/                             # patched vLLM checkout with SSR changes

Most users should start with the top-level scripts and only enter vllm/ when working on the implementation itself.

Installation

The replication environment used for the experiments is named specdecsd. The captured package lists are provided for reference:

environment-specdecsd.yml
requirements-specdecsd.txt

On the original machine, activate the environment with:

export SPECDEC_ENV=/data1/ravi/miniconda3/envs/specdecsd
export PATH="$SPECDEC_ENV/bin:$PATH"
export PYTHON="$SPECDEC_ENV/bin/python"

Set cache and temporary directories under /data1/ravi:

export HOME=/data1/ravi
export XDG_CACHE_HOME=/data1/ravi/.cache
export HF_HOME=/data1/ravi/.cache/huggingface
export HUGGINGFACE_HUB_CACHE=/data1/ravi/.cache/huggingface/hub
export TRANSFORMERS_CACHE=/data1/ravi/.cache/huggingface/hub
export TMPDIR=/data1/ravi/tmp
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export HF_HUB_DISABLE_XET=1
mkdir -p "$TMPDIR" "$HF_HOME"

Install the patched vLLM tree in editable mode:

cd /data1/ravi/code/specdec/specdec/vllm
VLLM_VERSION_OVERRIDE=0.11.2.dev278 "$PYTHON" -m pip install -e .

Check that Python imports vLLM from this repository:

cd /data1/ravi/code/specdec/specdec
"$PYTHON" - <<'PY'
import pathlib
import vllm
print(pathlib.Path(vllm.__file__).resolve())
PY

Expected path:

/data1/ravi/code/specdec/specdec/vllm/vllm/__init__.py

Data Preparation

The benchmark scripts default to:

/data1/ravi/specdec_data

Prepare HumanEval, ClassEval, LongProc, and the derived prompt files with:

cd /data1/ravi/code/specdec/specdec

PYTHON=/data1/ravi/miniconda3/envs/specdecsd/bin/python \
DATA_ROOT=/data1/ravi/specdec_data \
bash scripts/prepare_replication_data.sh

Generate the Qwen and Gemma source outputs used by the HumanEval/ClassEval latency runs:

PYTHON=/data1/ravi/miniconda3/envs/specdecsd/bin/python \
DATA_ROOT=/data1/ravi/specdec_data \
GPUS=0,1,2,3,4,5,6,7,8,9 \
bash scripts/prepare_replication_data.sh --generate-model-outputs

For a fast smoke test, use a separate data root and generate only tiny model outputs:

PYTHON=/data1/ravi/miniconda3/envs/specdecsd/bin/python \
DATA_ROOT=/data1/ravi/specdec_data_smoke \
GPUS=0 \
bash scripts/prepare_replication_data.sh --smoke --generate-model-outputs

More detail on data sources, normalization, and expected files is in data/REPLICATION_DATA.md.

Reproduction

The two result groups most worth reproducing are:

  • The full end-to-end latency table across Qwen3.5-4B and Gemma-4-E4B-it on LongProc 2K, ClassEval, and HumanEval.
  • The Gemma-4-E4B-it ClassEval ablation separating prefix verification and suffix decoding.

Full Code Benchmarks

Use this wrapper for the HumanEval/ClassEval rows from the main table and the Gemma ClassEval prefix/suffix ablation:

cd /data1/ravi/code/specdec/specdec

DATA_DIR=/data1/ravi/specdec_data \
GPUS=0,1,2,3,4,5,6,7,8,9 \
bash scripts/benchmarks/run_paper_code_latency_rows.sh

This script writes results under $DATA_DIR/paper_code_latency_rows/<run_tag>/.

The paper-row code settings are fixed-length settings:

CoT cap: 2000 tokens
Draft budget: 500 CoT tokens
Max draft length: 500 tokens
HumanEval continuation cap: 1500 tokens
ClassEval continuation cap: 4000 tokens
Suffix source: draft_only
Suffix tree depth / speculative tokens / factor: 4 / 16 / 4
Internal suffix cache: disabled

Do not use the source-length exploratory wrapper scripts/benchmarks/run_full_code_eval_suffix_latency.sh to reproduce the main latency table. That older wrapper uses COT_LENGTH_SOURCE=source, CONTINUATION_LENGTH_SOURCE=answer_minus_manual, and a smaller d2/s8/f2 suffix setting, which is useful for diagnostics but does not match the reported paper rows.

To run only one row, set EXPERIMENTS:

EXPERIMENTS=qwen_humaneval_main \
DATA_DIR=/data1/ravi/specdec_data \
GPUS=0,1,2,3,4,5,6,7,8,9 \
bash scripts/benchmarks/run_paper_code_latency_rows.sh

Supported code-row experiment names:

qwen_humaneval_main
gemma_humaneval_main
qwen_classeval_main
gemma_classeval_prefix_only
gemma_classeval_suffix_only
gemma_classeval_prefix_suffix

LongProc 2K

Use the full-prompt LongProc runner for the n=989 rows:

cd /data1/ravi/code/specdec/specdec

RUN_TAG=paper_longproc2k_full_repl_$(date +%Y%m%d_%H%M%S) \
PHASES=rejection \
QWEN_REJECTION_CONFIG=c1000_b500_m500_single \
QWEN_REJECTION_SUFFIX_MODE=suffix \
GEMMA_REJECTION_CONFIG=c2000_b500_m500_multi750 \
GEMMA_REJECTION_SUFFIX_MODE=suffix \
QWEN_REJECTION_GPUS=0,1,2,3,4 \
GEMMA_REJECTION_GPUS=5,6,7,8,9 \
DATA_DIR=/data1/ravi/specdec_data \
RESULTS_DIR=/data1/ravi/specdec_data/longproc2k_full_prompt_eval \
LOG_DIR=/data1/ravi/specdec_data/longproc2k_full_prompt_eval_logs \
RUNS=1 WARMUP_RUNS=0 \
bash scripts/benchmarks/run_longproc2k_full_prompt_eval.sh

LongProc 2K uses these five subsets:

countdown_2k
html_to_tsv_2k
path_traversal_2k
pseudo_to_code_2k
tom_tracking_2k

Gemma ClassEval Prefix/Suffix Ablation

The ablation uses the same benchmark runner with only the reuse mode changed:

prefix-only:
  SELF_SPEC_MODE=rejection
  no suffix-decoding flags

suffix-only:
  SELF_SPEC_MODE=manual
  MANUAL_ACCEPTED_PREFIX_LENGTH=0
  suffix-decoding flags enabled

prefix + suffix:
  SELF_SPEC_MODE=rejection
  suffix-decoding flags enabled

The reusable grid runner is:

scripts/benchmarks/bench_self_spec_real_tasks_grid.sh

The underlying Python benchmark is:

scripts/benchmarks/bench_vllm_self_spec_real_tasks.py

See notes/experiment_plan.md for the full experiment checklist and expected runtime order.

Summarizing Results

Summarize HumanEval/ClassEval JSONL files:

python scripts/analysis/summarize_self_spec_real_tasks.py \
  --output-md notes/experiment_results.md \
  --results-glob '/data1/ravi/specdec_data/self_spec_real_tasks*.jsonl'

Summarize LongProc 2K full-prompt results:

python scripts/analysis/summarize_longproc2k_full_eval.py \
  --results-glob '/data1/ravi/specdec_data/longproc2k_full_prompt_eval/self_spec_real_tasks.longproc2k_full_prompt.*.jsonl' \
  --output-md notes/longproc_full2k_eval_results.md

Known Local Replication

The following local checks were completed while preparing this release:

  • Data preparation smoke test under /data1/ravi/specdec_data_release_smoke.
  • Qwen and Gemma one-token model-output smoke test.
  • Byte-for-byte comparison between regenerated prompt files and the existing /data1/ravi/specdec_data prompt artifacts.
  • Gemma-4-E4B-it ClassEval prefix + suffix replication: 100 samples, 1.324x mean speedup, 24.5% latency reduction.
  • Gemma-4-E4B-it ClassEval prefix-only replication: 100 samples, 1.092x mean speedup, 8.4% latency reduction.
  • Qwen3.5-4B HumanEval corrected fixed-cap replication: 164 samples, 1.036x mean speedup, 3.50% latency reduction, using fixed cot=2000, fixed continuation=1500, and d4/s16/f4 draft-only suffix decoding.

More local run paths and numbers are recorded in notes/experiment_results.md.

Implementation Notes

The vLLM changes are concentrated in the vendored vllm/ tree. At a high level, the scheduler creates child draft requests tied to the original reasoning request, cancels stale draft children when the parent finishes thinking, and passes completed drafts into verification and suffix-backed continuation.

For implementation details, start with:

vllm/
scripts/benchmarks/bench_vllm_self_spec_real_tasks.py
scripts/benchmarks/bench_self_spec_real_tasks_grid.sh

For readers interested in implementing SSR themselves, the project began with a smaller prototype in vllm/nano-vllm/. That version is easier to read than the full vLLM integration and is useful for understanding the request flow before working through the production-style scheduler changes. In particular, see:

vllm/nano-vllm/SELF_SPEC_CHANGE_MAP.md
vllm/nano-vllm/test_qwen3_thinking_self_spec.py
vllm/nano-vllm/test_qwen3_thinking_multi_draft.py

Troubleshooting

If import vllm resolves to a system package instead of this repository, reinstall the vendored tree with pip install -e ./vllm and check PYTHON and PATH.

If downloads or model loads write outside /data1/ravi, check HF_HOME, HUGGINGFACE_HUB_CACHE, TRANSFORMERS_CACHE, and TMPDIR.

If multiprocessing fails inside vLLM workers, set:

export VLLM_WORKER_MULTIPROC_METHOD=spawn

If Hugging Face downloads use Xet unexpectedly, set:

export HF_HUB_DISABLE_XET=1

Acknowledgements

This work builds on vLLM and its V1 suffix decoding implementation. The experiments use HumanEval, ClassEval, and LongProc. The method is related to speculative decoding, self-speculative decoding, prompt lookup and suffix decoding, and inference-time reasoning budget control.

Citation

If you use this codebase, please cite the accompanying paper once it is released:

@misc{selfspec2026,
  title = {Self-Speculation for Faster Reasoning Models},
  year = {2026},
  note = {Citation details to be updated at release}
}

About

No description, website, or topics provided.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages