TorchSpec

TorchSpec is a torch-native speculative decoding training framework. We introduce a disaggregated way of training speculative decoding draft models where inference and training are fully decoupled and stream hidden states directly from inference engine groups to distributed training workers via Mooncake store, allowing each side to scale independently.

Adoption

TorchSpec has been adopted by production inference platforms and the vLLM ecosystem:

• DigitalOcean used TorchSpec to train an EAGLE3 draft model for MiniMax-M2.5 on DigitalOcean Serverless Inference.
• vLLM used TorchSpec and vLLM to train the custom EAGLE3 draft model featured in its Artificial Analysis leaderboard work.
• CoreWeave used TorchSpec to train a DFlash speculative decoding model for Kimi K2.7 Code and contributed D-PACE support upstream.
• fal used TorchSpec to train a DSpark speculative decoding model for its Ideogram V4 prompt expander, reporting 16x throughput gains with DSpark.
• Tencent Hunyuan open-sourced AngelSpec, adopting TorchSpec for training and featuring it in the paper.

🤗 Released Models

Draft models trained with TorchSpec, available on the LightSeek Foundation Hugging Face organization:

• lightseekorg/kimi-k2.5-eagle3
• lightseekorg/kimi-k2.5-eagle3-mla
• lightseekorg/kimi-k2.6-eagle3
• lightseekorg/kimi-k2.6-eagle3-mla
• lightseekorg/kimi-k2.6-eagle3.1-mla
• lightseekorg/kimi-k3-eagle3-mla
• lightseekorg/kimi-k3-dflash2
• lightseekorg/kimi-k3-dspark

Draft models trained with TorchSpec, available from other organizations:

• Inferact/MiniMax-M3-EAGLE3
• Inferact/Kimi-K3-DSpark

🚀 Blogs

Blogs and Announcements

• Releasing Kimi K3 Draft Collection: Efficient Draft Training in TorchSpec (Sep 2026)
• EAGLE 3.1: Advancing Speculative Decoding Through Collaboration Between the EAGLE Team, vLLM, and TorchSpec (May 2026)
• TorchSpec: Speculative Decoding Training at Scale (PyTorch Blog, March 2026)
• TorchSpec: Speculative Decoding Training at Scale (March 2026)

Table of Contents

• Architecture Overview
• Inference Backend Support
• Quick Start
• Setup
• Examples
• Training Modes
• Checkpoint Conversion
• Metrics Reporting
• Troubleshooting

Architecture Overview

TorchSpec is built around a disaggregated training pipeline:

• Inference engines generate target-model hidden states with inference engines.
• Mooncake store transfers tensors between inference and training without materializing them on disk.
• Training workers consume streamed hidden states to train speculative decoding draft models.

This separation keeps the training side focused on optimization while letting the inference side scale for hidden-state generation throughput.

Inference Backend Support

TorchSpec streams hidden states from inference engines into training workers.

Quick Start

Train an Eagle3 draft model for Qwen3-8B on a single node with 4 GPUs (2 for training and 2 for inference):

./examples/qwen3-8b-single-node/run.sh

Override config values directly from the CLI:

./examples/qwen3-8b-single-node/run.sh training.learning_rate=5e-5 training.num_train_steps=500

Setup

Quick Setup

# Install with vLLM
./tools/build_conda.sh 1 vllm
micromamba activate torchspec

# Or install with SGLang
./tools/build_conda.sh
micromamba activate torchspec

# Or install TokenSpeed from an editable source checkout
./tools/build_conda.sh 1 tokenspeed
micromamba activate torchspec

To install into your current environment instead:

./tools/build_conda.sh current tokenspeed  # or 'sglang', 'vllm', or 'both'

The TokenSpeed backend currently requires a Python 3.12 environment because its native kernel dependency wheels do not support Python 3.14.

Optional: install Flash Attention support:

pip install -e ".[fa]"

Backend-Specific Usage

vLLM

./examples/qwen3-8b-single-node/run.sh configs/vllm_qwen3_8b.yaml

SGLang

./examples/qwen3-8b-single-node/run.sh

TensorRT-LLM

Run inside the TensorRT-LLM image (docker/trtllm/v1.3.0rc18/Dockerfile), which ships tensorrt_llm pre-patched for Mooncake hidden-state capture:

./examples/qwen3-8b-single-node/run.sh configs/trtllm_qwen3_8b.yaml

Single-node tensor parallelism only for now (multi-node TP is not yet wired up).

TorchSpec uses vLLM's Worker Extension mechanism to hook into the model forward pass and capture hidden states directly inside worker processes, which avoids RPC serialization overhead during extraction. For SGLang, TorchSpec applies a patch to the existing codebase to enable hidden-state extraction. For TensorRT-LLM, TorchSpec builds on its native SaveHiddenStates speculative mode and applies a small patch that redirects the captured aux + final hidden states to Mooncake instead of writing them to disk.

Updating the SGLang patch

Commit your changes in the _sglang checkout, then regenerate patches/sglang/<version>/sglang.patch from the pinned base commit with ./tools/update_sglang_patch.sh.

Examples

See examples/README.md for more details about each example.

Training Modes

Offline Replay Training

Offline replay (Docs) is a training mode that separates target-output generation from draft-model training by reading hidden states from disk. Recommended for testing and development on 1 GPU.

Resume vs. Continual Training

Both modes use training.load_path, but they restore different states:

Resume the same run:

training:
  load_path: /path/to/old_run/checkpoints

output_dir: /path/to/old_run

Start a new run from existing weights:

training:
  load_path: /path/to/old_run/checkpoints
  continual_training: true
  learning_rate: 1e-5
  warmup_ratio: 0.01
  num_epochs: 1

output_dir: /path/to/new_run

Start a new run from an existing Hugging Face checkpoint. --input accepts a Hugging Face Hub repo id (downloaded automatically), a local HF/safetensors directory or .safetensors file, or a TorchSpec DCP checkpoint dir:

python tools/convert_to_torchspec.py \
    --input org/dflash-checkpoint \
    --config torchspec/config/dspark_draft_config_qwen36_35b.json \
    --output ./outputs/dspark_init

Then warm-start from the generated init:

training:
  load_path: ./outputs/dspark_init
  continual_training: true

Using this technique, you can also warm-start DSpark training runs from a pre-trained DFlash model for faster convergence.

Checkpoint Conversion

Convert an FSDP checkpoint to HuggingFace format:

python tools/convert_to_hf.py --input-dir ./outputs/my_experiment/iter_0010000/

Vocabulary pruning, which reduces the draft model lm_head to a smaller token set and emits d2t and t2d mappings, can be applied either during training or at conversion time.

DFlash2 requires the full target vocabulary. Do not set draft_vocab_size or pass --prune-vocab for a DFlash2 model.

• Pre-pruning: set draft_vocab_size in your training config. The checkpoint already contains the pruned lm_head and d2t/t2d buffers, so the basic conversion command is enough.
• Post-pruning: train with the full vocabulary, then pass --prune-vocab at conversion time together with a representative dataset to compute token frequencies.

python tools/convert_to_hf.py \
    --input-dir ./outputs/my_experiment/iter_0010000/ \
    --prune-vocab \
    --dataset-path Aeala/ShareGPT_Vicuna_unfiltered \
    --draft-vocab-size 32000 \
    --tokenizer Qwen/Qwen3-8B \
    --chat-template qwen \
    --prompt-key conversations

Pass --cache-dir ./cache to reuse the tokenized dataset cache from training.

Metrics Reporting

W&B logging is disabled by default with report_to: none. To enable it, set report_to: wandb in your config and provide your API key.

Troubleshooting

Set TORCHSPEC_LOG_LEVEL=DEBUG for more verbose logging when diagnosing issues:

TORCHSPEC_LOG_LEVEL=DEBUG ./examples/qwen3-8b-single-node/run.sh

Mooncake SEGFAULT

Current Mooncake version has a bug with TCP-only hosts causing a SEGFAULT error. Set MC_STORE_MEMCPY=0 until the upstream issue is fixed.

RDMA failure

If you get the error ... libcudart symbols not found globally. Make sure PyTorch with CUDA is installed before using TileLang, reinstall mooncake to match your system CUDA version, i.e. if you run on CUDA 13, mooncake is probably installed for cu12 wheels, re-install it using uv pip uninstall mooncake-transfer-engine && uv pip install mooncake-transfer-engine-cuda13.

Per-Rank File Logging

Set TORCHSPEC_LOG_DIR to an absolute path on a shared filesystem (NFS) to enable per-rank log files for every Ray actor on both training and inference:

export TORCHSPEC_LOG_DIR=/my_project/running_logs

This creates a structured directory with one file per actor, organized by role and node:

running_logs/
  training/
    10.0.0.1/
      training_g0_rank0_20260301_080012.log
      training_g0_rank1_20260301_080012.log
    10.0.0.2/
      training_g0_rank2_20260301_080013.log
  inference/
    10.0.0.1/
      inference_g0_rank0_20260301_080014.log
    10.0.0.2/
      inference_g0_rank1_20260301_080015.log

The path must be absolute and writable from all nodes. If TORCHSPEC_LOG_DIR is unset or not writable, per-rank file logging stays disabled and Ray falls back to stdout/stderr capture.