specdecode-rl

A small trace-replay environment for adaptive speculative decoding.

At every verification round, a policy picks k ∈ {0, 1, 2, 4, 8, 16}. Larger k can commit more tokens, but drafting and verification both cost time. k=0 skips speculation and runs the target model normally.

I built this to see whether an RL policy could learn that trade-off from real GPU traces. Building the environment means deciding what the agent can see, what it controls, and what it should optimize. In this case, acceptance rate alone is a bad reward. Latency is what matters.

Prime Intellect environment · W&B runs

Setup

The traces come from 50 prompts in nvidia/SPEED-Bench, profiled on an A100-SXM4-80GB.

Target: Qwen/Qwen3-8B
Draft:  Qwen/Qwen3-1.7B
Dtype:  bfloat16
Output: 48 tokens
Split:  35 train / 5 validation / 10 test
Actions: 0, 1, 2, 4, 8, 16
Reward: 1 - policy_latency / target_only_latency

The simulator replays the measured draft and verification costs. A hindsight oracle uses dynamic programming to find the cheapest action at each token position.

Results

Mean speedup on the same 10 held-out traces, with three seeds for learned policies:

Policy	Speedup
Target-only	1.0000×
REINFORCE	1.0056× ± 0.0145
Fixed k=2	1.0192×
Oracle imitation	1.0393× ± 0.0135
Hindsight oracle	1.2307×

Vanilla REINFORCE was mostly a miss. It collapsed toward near-fixed actions and failed to beat fixed k=2. Training the same policy to imitate oracle actions worked better: it learned to speculate for a few rounds, then fall back to target-only decoding.

I also removed groups of observations to check what the policy was using:

Features	Speedup
All features	1.0393× ± 0.0135
No explicit token position	1.0362× ± 0.0123
No acceptance/confidence history	1.0283× ± 0.0183
Acceptance/confidence only	1.0180× ± 0.0000

Previous acceptance and draft confidence help. Explicit token position barely matters, although elapsed latency still acts as a rough progress signal.

Full experiment notes:

• experiments/2026-08-15-qwen3-reinforce.md
• experiments/2026-08-15-qwen3-oracle-imitation.md

Run

uv sync --extra dev
uv run pytest

uv run python scripts/evaluate_traces.py \
  --trace-path traces/speed-qwen3-8b-1.7b.jsonl

Train an imitation policy locally:

from specdecode_rl.imitation import ImitationConfig, train_imitation

result = train_imitation(ImitationConfig(
    trace_path="traces/speed-qwen3-8b-1.7b.jsonl",
    checkpoint_path="checkpoints/policy.pt",
))
print(result.test_eval.mean_speedup)

Profile new traces on an NVIDIA GPU:

uv sync --extra profile
uv run specdecode-profile \
  --target-model Qwen/Qwen3-8B \
  --draft-model Qwen/Qwen3-1.7B \
  --output traces/speed-qwen3.jsonl \
  --limit 8

Profiling writes a manifest and supports --resume.

Caveat

This is offline, deterministic, batch-size-1 trace replay. It is not wired into vLLM or SGLang, and the numbers are not production serving speedups. The dataset is also small, and the test traces had appeared in earlier aggregate exploration before being held out from training.