Skip to content

Latest commit

 

History

6 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

specdec

Measure the speculative-decoding acceptance rate and expected speedup for a target/draft model pair — before you spend a GPU-month training a draft.

Acceptance rate is the objective every draft-training method (GKD, Medusa, EAGLE-3, …) actually optimises:

alpha  =  E_x[ sum_v min( p_target(v | x), p_draft(v | x) ) ]  =  1 - TV(p_target, p_draft)

specdec computes it directly, so you can tell a useful draft from a useless one, sanity -check a training run, or reproduce a paper's number — all with plain forward passes that run on a laptop.

Install

pip install -e .            # core: torch + transformers
pip install -e '.[data]'    # + datasets, for `--data hf:...`
pip install -e '.[dev]'     # + pytest

Use

specdec measure \
  --target Qwen/Qwen2.5-1.5B-Instruct \
  --draft  Qwen/Qwen2.5-0.5B-Instruct \
  --data builtin \
  --temperature 0.7 --gamma 4 \
  --continuation-tokens 64 \
  --simulate
  speculative-decoding acceptance  —  Qwen/Qwen2.5-0.5B-Instruct  ⇒  Qwen/Qwen2.5-1.5B-Instruct
  sampling: temperature=0.7    gamma (draft block): 4
  ──────────────────────────────────────────────────────────────
  distributional (teacher-forced on target continuation)
    acceptance rate  alpha        0.781   ± 0.061 (per-seq)
    scored positions             2,048 over 32 sequences
    E[accepted draft tokens]      2.341  of 4
    E[tokens per cycle]           3.341
    est. speedup                   2.09x   (cost_ratio=0.15, rough)
  ──────────────────────────────────────────────────────────────
  generative (real rejection sampling — the number to trust)
    acceptance rate  alpha        0.763
    mean accepted length          2.218  of 4
    tokens per cycle              3.218
    est. speedup                   2.01x   (cost_ratio=0.15, rough)
    alpha by draft offset 1..4  0.84  0.78  0.73  0.69
      (a steep fall = the draft degrades as it runs ahead)

Use --continuation-tokens. Without it, acceptance is scored on the prompt — human -written question text. At serving time the draft is predicting the target's own output, and acceptance on the two differs substantially. --continuation-tokens 64 samples a continuation from the target and scores only those positions.

As a library

from specdec import measure_acceptance, SamplingParams
from specdec.models import load_model_and_tokenizer
from specdec.data import build_sequences, load_texts

target, tok, dev = load_model_and_tokenizer("Qwen/Qwen2.5-1.5B-Instruct")
draft,  _,  _   = load_model_and_tokenizer("Qwen/Qwen2.5-0.5B-Instruct", device=dev)

seqs = build_sequences(load_texts("builtin"), tok, max_samples=16)
report = measure_acceptance(
    target, draft, seqs,
    sampling=SamplingParams(temperature=0.7), device=dev,
    continuation_tokens=64,   # score the target's own output, not the prompt
    simulate=True,            # the number to trust
)
print(report.alpha)                    # distributional
print(report.empirical_alpha)          # generative
print(report.alpha_by_draft_offset)    # degradation curve, simulation only

What the numbers mean, and what they don't

  • alpha (distributional) is measured teacher-forced over each scored position: the exact expected one-step acceptance if the draft proposed from that same context. It is the standard "acceptance rate" reported in the literature, and it is cheap — two forward passes, no generation loop.

  • --simulate (generative) runs the real Leviathan/Chen rejection-sampling loop: the draft proposes from its own output, tokens are accepted/rejected exactly, and rejections resample from the normalised residual (p-q)+. Trust this one. It is the only mode that captures the draft conditioning on its own proposals.

  • alpha by draft offset is P(accept at offset j | offset j was reached) and comes only from the simulation. Read it with the cycles reaching it row underneath. It is a conditional probability on a shrinking, self-selected cohort, and two effects fight:

    • drift pushes it down — by offset j the draft is conditioning on j of its own tokens and has wandered further from the target. This is what EAGLE-3's training-time test exists to fix.
    • selection pushes it up — only cycles where the draft already agreed j times reach offset j, and those are the easy, predictable contexts.

    On degenerate output selection can win outright. Greedy GPT-2 falls into repetition loops, and the tail rises to 0.95 on 74 surviving cycles: once the draft is inside the loop, agreeing again is nearly free. A - means the offset was never reached at all — itself a strong signal that the draft is rejected before it gets there.

    (None of this is measurable teacher-forced: with every position conditioned on the reference context there is no notion of "draft offset" at all.)

  • est_speedup is tokens_per_cycle / (gamma * cost_ratio + 1). It ignores KV-cache warmup, batching, and memory bandwidth. It is an indication, not a benchmark — measure real latency on real hardware before quoting a number.

  • Temperature matters. --temperature 0 gives exact top-1 agreement, which over -estimates what a sampling deployment sees. Always measure at your serving temperature.

  • Standard speculative decoding needs a shared tokenizer. specdec truncates logits to the common vocab prefix (before the draft samples, so it can never propose a token the target lacks) and warns; results are only exact when the surplus tokens are padding.

  • Precision defaults to fp32 off CUDA, as the conservative choice for a tool whose whole output is a comparison of two distributions. The cost of not doing so turns out to be small — logits are cast to fp32 before any softmax, so half precision only perturbs the forward pass. Measured on gpt2 ⇐ distilgpt2, fp16 moves α by 0.001 (0.71113 → 0.71216) against a per-sequence spread of 0.084. Use --dtype float16 when memory matters: a 1.5B target plus a 0.5B draft will not fit in 16GB at fp32.

GPU / Unsloth

--use-unsloth loads via unsloth.FastLanguageModel (needs CUDA/ROCm/XPU). The default plain-transformers path runs anywhere, including Apple Silicon (mps).

Validation

pytest            # 64 tests, no model downloads; pure-tensor maths + stand-in models

Beyond the unit tests, the two measurement paths are checked against each other. They are produced by completely independent routes — a closed-form Σ min(p,q) over teacher-forced logits, versus an actual rejection-sampling loop with an RNG — so agreement is meaningful evidence that both are right.

gpt2 ⇐ distilgpt2, a pair with a known relationship (distilgpt2 was distilled from gpt2), 8 prompts / 512 generated tokens:

sampling distributional α generative α agreement
--temperature 0.7 0.711 0.715 0.004
--temperature 0 (greedy) 0.840 0.835 0.005

Greedy exceeding sampled (0.84 > 0.71) is the expected direction: top-1 agreement over-estimates what a sampling deployment sees.

Qwen2.5-1.5B-Instruct ⇐ Qwen2.5-0.5B-Instruct, a current same-family pair, chat template, --temperature 0.7:

run distributional α generative α
6 prompts, 64-token continuation 0.725 ± 0.099
4 prompts, 48-token continuation, --simulate 0.775 ± 0.043 0.724

α ≈ 0.72–0.78 is where the literature puts a same-family draft of this size ratio. The 0.051 gap in the second row is 1.19σ given 145 examined draft tokens and n=4 sequences — sampling noise, not a discrepancy. The gpt2 rows are the tighter evidence because they carry 2–4× the sample.

Endpoint checks: identical models give α = 1.000 and 5.000 tokens/cycle at γ=4; an untrained draft (sshleifer/tiny-gpt2) against distilgpt2 gives α ≈ 0.01 and an estimated speedup of 0.63× — correctly reporting that the draft would make you slower.

Reproduce:

specdec measure --target gpt2 --draft distilgpt2 \
  --data builtin --max-samples 8 --max-length 40 \
  --temperature 0.7 --gamma 4 --continuation-tokens 64 --simulate

License

Apache-2.0.

About

Measure speculative-decoding acceptance rate and expected speedup for a target/draft model pair

Topics

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages