[Feature]: Add speculative decoding with draft model pruning

[jmamou]

🚀 The feature, motivation and pitch

I have implemented the FR-Spec approach
at the logits processor level, using AllowedTokenIdsLogitsProcessor. This implementation does not prune the draft model itself but allows evaluating acceptance rates under different draft pruning ratios. You can find the code here.

MT-Bench results:

pruning ratio	vanilla	0.1	0.25	0.5	0.75	0.9	0.99
draft acceptance rate (%)	27.8	28.3	28.6	28.6	27.2	25.9	18.8

New speculative config parameters:

• token_ids_by_frequency: Path to a tensor file containing token frequencies sorted by token IDs, used for pruning-based speculative decoding.
• pruning_ratio: Ratio of tokens to prune during speculative decoding.

Example usage

VLLM_USE_V1=1 vllm serve meta-llama/Llama-3.1-8B-Instruct \
  --speculative_config '{"method": "eagle", "model": "yuhuili/EAGLE-LLaMA3.1-Instruct-8B", "num_speculative_tokens": 4, "token_ids_by_frequency": "vllm/examples/target-dist-meta-llama-Llama-3.1-8B-Instruct-wikitext-wikitext-103-raw-v1-train.pt", "pruning_ratio": 0.1}'

vllm bench serve \
  --model meta-llama/Llama-3.1-8B-Instruct \
  --endpoint-type openai-chat \
  --endpoint /v1/chat/completions \
  --dataset-name hf \
  --dataset-path philschmid/mt-bench \
  --num-prompts 80 \
  --max-concurrency 16 \
  --temperature 1 \
  --top-p 1.0

By selectively pruning unlikely tokens in the draft model, this feature is expected to improve speculative decoding speedups while maintaining high acceptance rates, enabling faster inference for large models like LLaMA-3.1-8B-Instruct.

@keyboardAnt @eitanturok
FR-Spec implementation #24343

Alternatives

No response

Additional context

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[keyboardAnt]

Thank you, @jmamou.

MT-Bench results:

pruning ratio vanilla 0.1 0.25 0.5 0.75 0.9 0.99
draft acceptance rate (%) 27.8 28.3 28.6 28.6 27.2 25.9 18.8

How come acceptance of vanilla (pruning rate = 0) isn't maximal? (Over how many repeats are the results averaged?)

[jmamou]

Thank you, @jmamou.

MT-Bench results:
pruning ratio vanilla 0.1 0.25 0.5 0.75 0.9 0.99
draft acceptance rate (%) 27.8 28.3 28.6 28.6 27.2 25.9 18.8

How come acceptance of vanilla (pruning rate = 0) isn't maximal? (Over how many repeats are the results averaged?)

It seems insignificant, 80 samples in a single run

[benchislett]

@jmamou with a temperature of 1.0, there might be some variance. Could you run with a lower temperature, and/or average over multiple runs? It does seem odd that there would be that much variance, so worth confirming.

[eitanturok]

Great work @jmamou!

Comparing my implementation to yours, I think we'd see similar acceptance lengths but different speedups.

Your approach computes logits for the full vocabulary first, then prunes - but we can get better performance by pruning the vocabulary before computing logits. That's the key difference in my implementation.

[keyboardAnt]

Thank you, @jmamou.

MT-Bench results:
pruning ratio vanilla 0.1 0.25 0.5 0.75 0.9 0.99
draft acceptance rate (%) 27.8 28.3 28.6 28.6 27.2 25.9 18.8

How come acceptance of vanilla (pruning rate = 0) isn't maximal? (Over how many repeats are the results averaged?)

It seems insignificant, 80 samples in a single run

Are the different pruning ratios compared using the same initial prompts? Same random seed?

[keyboardAnt]

Comparing my implementation to yours, I think we'd see similar acceptance lengths but different speedups.

@eitanturok - great! Can you please share your measured acceptance in #24343?

[jmamou]

Great work @jmamou!

Comparing my implementation to yours, I think we'd see similar acceptance lengths but different speedups.

Your approach computes logits for the full vocabulary first, then prunes - but we can get better performance by pruning the vocabulary before computing logits. That's the key difference in my implementation.

@eitanturok

Right, my implementation doesn’t provide any speedup since the draft wasn’t pruned. The goal of my experiments was only to verify that the acceptance rate remains nearly unchanged when pruning the logits of less frequent tokens.

[jmamou]

Thank you, @jmamou.

MT-Bench results:
pruning ratio vanilla 0.1 0.25 0.5 0.75 0.9 0.99
draft acceptance rate (%) 27.8 28.3 28.6 28.6 27.2 25.9 18.8

How come acceptance of vanilla (pruning rate = 0) isn't maximal? (Over how many repeats are the results averaged?)

It seems insignificant, 80 samples in a single run

Are the different pruning ratios compared using the same initial prompts? Same random seed?

same prompts.
I am not sure about the random seed.

[jmamou]

Additional results running spec_decode.py

VLLM_USE_V1=1 python3 examples/offline_inference/spec_decode.py \
                    --method eagle \
                    --num_spec_tokens 4 \
                    --dataset-name hf \
                    --dataset-path philschmid/mt-bench \
                    --num_prompts 100 \
                    --model-dir 'meta-llama/Llama-3.1-8B-Instruct' \
                    --eagle-dir 'yuhuili/EAGLE-LLaMA3.1-Instruct-8B'

	pruning ratio	mean acceptance length (temp=0)	mean acceptance length (temp=0.9)
vanilla		2.39	2.29
fspec	0.25	2.37	2.27
frspec	0.5	2.34	2.22
frspec	0.75	2.28	2.2
frspec	0.9	2.18	2.1
frspec	0.99	1.84	1.8

[keyboardAnt]

The following micro-benchmark demonstrates the advantage of vocabulary pruning. The plot shows the median latency (ms) over the batch size for a fixed hidden dimension of 1024. The speedup becomes more pronounced as the batch size increases. This suggests that vocabulary pruning methods like frspec will work best for larger batch sizes.

Hardware: A100
Code: https://github.com/keyboardAnt/vllm/blob/nadav/bench-vocab-pruning/benchmarks/kernels/bench_vocab_pruning.py (commit 4ef7eb2)

[keyboardAnt]

We compute arithmetic intensity (AI) as FLOPs per byte of data moved. For a matmul [B, H] × [H, V] → [B, V]:

• FLOPs: 2 * B * H * V (each multiply-accumulate = 2 FLOPs)
• Bytes (weights + activations + output logits): dtype_size * (H*V + B*H + B*V) (with dtype_size = 2 bytes for bf16)

So: AI = (2 * B * H * V) / (dtype_size * (H*V + B*H + B*V))

We report the including-outputs version (roofline-style).

Example (B=1024, H=1024, V=32000):

FLOPs = 2 * 1024 * 1024 * 32000 ≈ 67.1e9
Bytes = 2 * (1024*32000 + 1024*1024 + 1024*32000) ≈ 132e6
AI    ≈ 510 FLOPs/byte

(Same hardware and code as in my previous comment)

There is also a variant that excludes output bytes: AI_excl_output = (2 * B * H * V) / (dtype_size * (H*V + B*H)). This is sometimes used for fused kernels where logits do not hit HBM.