TorchSpec: Speculative Decoding Training at Scale

TL;DR

TorchSpec is a torch-native framework for training speculative-decoding draft models (EAGLE-3-style) at scale, by fully disaggregating the inference system that produces target-model hidden states from the training system that consumes them. Hidden states are streamed via the Mooncake transfer engine (RDMA/TCP, GPU-direct, zero-copy) instead of being precomputed to disk or co-located on the same GPUs. With this, the team trained a Kimi K2.5 EAGLE-3 draft model on 1500 H200-hours over 600k samples (6B tokens), removing the storage cost that would otherwise be ~700 TB for a 100k-sample dataset and freeing the training GPUs of the 575 GB target-model weight so context lengths can reach 100k–200k tokens.

Key claims

Standard speculative-decoding draft training has two unsatisfying options: precompute hidden states to disk (e.g. ~7 GB per 128K-token sample for Kimi K2.5 ⇒ 70 TB / 10k samples) or co-locate target + draft on the same GPUs (target eats ~575 GB, leaving ~8 GB/GPU on 8×H100, capping context length at ~4096 tokens) [§Background; “Storage analysis” and “Memory analysis” tables].
Disaggregation: target runs on a dedicated inference cluster, draft trains on a separate cluster, hidden states stream between them through a central Mooncake store via RDMA or TCP, so inference and training scale independently [§TorchSpec].
Mooncake provides RDMA/TCP with a unified API, GPU-Direct RDMA into pinned buffers, and zero-copy transfers — addressing the gigabytes-per-sample throughput requirement [§“Why Mooncake?”].
The inference path uses the production inference engine itself (vLLM, SGLang; TensorRT-LLM coming), so tokenization, kernels, and templates are bit-identical between training and deployment, eliminating train/serve skew for the draft model [§“Long Context Support”].
“Train with decode”: because hidden states are generated by a live inference engine rather than pre-tokenized files, prompt-only datasets can be expanded autoregressively at training time, avoiding the two-stage “regenerate responses with target, then train draft” pipeline [§“Long Context Support” / Train with Decode].
Reported deployment: 2 nodes of 8×H200 (TP=8 inference + TP=8 training) or 3 nodes of 8×H100 (TP=16 inference + TP=8 training) trains a Kimi K2.5 EAGLE-3 draft model over 600k samples (6B tokens) in 1500 H200-hours [§“Case Study”].
Resulting K2.5 draft model: output-throughput improvements of +60% at batch=1, +30% at batch=8, +26% at batch=16 with lookahead=3 over autoregressive Kimi K2.5 serving [§Introduction, draft-trained-with-lookahead=4 row].
A single B200 GPU can train an EAGLE-3 draft for Kimi K2.5 with sequence lengths up to 200k tokens; a single H100 reaches 44k — both impossible under the co-located approach (capped at ~4096 with 8 GB/GPU residual memory) [§“Long Context Support”].

Method

The pipeline has three components: (1) an inference engine group (vLLM/SGLang) holding the target model on dedicated GPUs, producing the 3 layers of intermediate hidden states required by EAGLE-3 (~7 GB per 128K-token sample for Kimi K2.5); (2) a Mooncake store acting as the cross-cluster transfer engine, written into via GPU-Direct RDMA from inference workers and read from by training workers, with no disk involved; (3) a training worker group running FSDP on the draft model with its full GPU memory available for activations (important for Training-Time Testing / TTT-style EAGLE-3 training, which retains activations across multiple speculative steps). Because the inference engines are the same code that serves the model in production, the draft model trains on exactly the hidden-state distribution it will see at deployment time. Two reference launch configurations are published: 2×(8×H200) with TP=8 on both sides, and 3×(8×H100) with TP=16 inference + TP=8 training. The trained draft model and the 600k-sample mixed dataset are both open-sourced on HuggingFace (lightseekorg/kimi-k2.5-eagle3, lightseekorg/kimi-mtp-dataset).

Results

End-to-end serving speedup with the trained draft model on Kimi K2.5: +60% throughput at batch=1, +30% at batch=8, +26% at batch=16 under lookahead=3 — i.e. the draft model trained via TorchSpec is “good enough” to justify the speculative-decoding overhead at multi-tenant batch sizes, not just at batch=1 [§Introduction].
Storage avoided: 700 TB for 100k samples / 210 TB for 30k samples on the offline-precompute alternative — eliminated entirely [§Background, “Storage analysis”].
Memory unlocked: co-located training on 8×H100 leaves ~8 GB/GPU residual, capping context at 4096; TorchSpec’s dedicated training cluster reaches 44k tokens on a single H100 and 200k on a single B200 [§“Long Context Support”].
Training cost: Kimi K2.5 draft model trained on 600k samples (~6B tokens) in 1500 H200-GPU-hours total [§“Case Study”].

Why it’s interesting

The result lives at the intersection of two recent threads in the wiki: lossless inference acceleration via speculative decoding (Speculative Speculative Decoding — Saguaro/SSD’s parallelization of draft and verify) and the broader push toward disaggregated LLM serving architectures (Mooncake itself originated for KV-cache transfer in Moonshot’s serving stack). Where Saguaro reorganizes draft↔verify scheduling at inference time to hide drafting latency, TorchSpec attacks the upstream problem of producing a high-quality draft model in the first place for trillion-parameter targets — the same bottleneck that has historically forced labs to either constrain context length (co-located) or burn 700 TB of disk (offline). Worth tracking because it’s also one of the first public showcases of a fully disaggregated PyTorch training pipeline where the inference engines are the production ones (vLLM/SGLang) rather than a separate trainer-side rewrite — a pattern that may generalize beyond speculative decoding (RL rollouts, distillation, reward-model inference for online RLHF). Contrast with veScale-FSDP: Flexible and High-Performance FSDP at Scale, which sits at the FSDP-internals layer of the stack; TorchSpec is the layer above that orchestrates two FSDP clusters communicating tensors over the network.