[V1][Spec Decode][Feature] Spec decode with probs

[andylolu2]

Summary

This PR enables speculative decoding with draft probabilities in V1. A partial revert of #16899.

Implementation

Blocker of #16899 was the draft probabilities aren't used immediately, so we need to keep them around for the next iteration. In this PR, I propose we add draft probabilities as part of the CachedRequestState. On the next iteration where the draft token ids of a request is used, we fetch the draft probs from the cached state. This greatly simplifies the matter:

• CachedRequestState already encapsulates the logic that data related to a request isn't necessarily used immediately. For example, this handles preemption / the problem where requests might not be scheduled on every step.
• Since the draft probs are tied to the cached state, they will be deallocated when the cached state is deleted as well, so little risk of memory leak (as opposed to managing a new cache).
• For small models, using draft probabilities is not worth the draft sampling/rejection sampling overhead. I've added a "kill-switch" enable_draft_probs in SpeculativeConfig for those cases. I recommend still enabling draft probs by default since it's robust to a wider range of temperatures (i.e. less likely to regress vs standard sampling).

Benchmark

Numbers obtained by running the follow on current main (14601f5) vs this branch:

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 \
    --temp <T>

Temperature	AL w/o probs	AL w/ probs
0	2.37	2.37
0.3	2.36	2.38
0.5	2.33	2.35
0.7	2.3	2.31
1	2.09	2.21
1.3	1.32	1.98
1.5	1.13	2.07
2	1.01	2.66

Old numbers w/ https://raw.githubusercontent.com/SafeAILab/EAGLE/49b67c8d2dc6d8154eab6ef6899e4b9cc4cd3e06/eagle/data/mt_bench/question.jsonl

Temperature	AL w/o probs	AL w/ probs
0	2.18	2.18
0.3	2.14	2.2
0.5	1.85	1.87
0.7	1.26	1.39
1	1.02	1.48
1.3	1	2
1.5	1	2.28
2	1	2.81

Old old numbers with https://raw.githubusercontent.com/SafeAILab/EAGLE/49b67c8d2dc6d8154eab6ef6899e4b9cc4cd3e06/eagle/data/mt_bench/question.jsonl

Temperature	AL w/o probs	AL w/ probs
0	2.29	2.29
0.3	2.1	2.12
0.5	1.79	1.84
0.7	1.26	1.3
1	1.02	1.47
1.3	1	1.98
1.5	1	2.27
2	1	2.8

Online benchmarks

Ran with:

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}'

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 <T> \
  --top-p 1.0

(I got some request-id-already-running error when using --num-prompts 100.)

Result

Temperature	TPOT w/o probs	TPOT w/ probs
0	4.98	5.03
0.3	5.2	5.44
0.5	5.21	5.3
0.7	5.37	5.28
1	6.06	5.66
1.3	9.07	6.29
1.5	10.61	5.77
2	11.62	4.58

Old numbers with https://raw.githubusercontent.com/SafeAILab/EAGLE/49b67c8d2dc6d8154eab6ef6899e4b9cc4cd3e06/eagle/data/mt_bench/question.jsonl

Temperature	TPOT w/o probs	TPOT w/ probs
0	4.79	4.78
0.3	4.91	5.34
0.5	5.42	5.47
0.7	5.8	5.75
1	6.92	6.13
1.3	9.96	6.24
1.5	11.38	5.62
2	11.86	4.53

Explanation:

• Overall acceptance rate is always higher after this PR, but only noticeable when temperature is sufficiently large (> 0.7).
• As temperature grows larger, acceptance rate also grows. This is because both the draft and target distribution is getting more and more uniform, hence more similar.
• Sampler overhead means the delta is not super large, but this can be optimized in a follow up.

Tests

I've mainly added two tests:

• test_propose_random that checks the draft model is sampling truthfully according to the underlying logits. I do so by sampling N times and ensure the observed distribution statistically don't deviate too much from the target.
• TestSpecDecodeScheduling which checks for edge cases like preemption and ensure that the draft probs are still cached properly and there's no things like index-OOB errors.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[gemini-code-assist]

Summary of Changes

Hello @andylolu2, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the speculative decoding feature, particularly for the Eagle model, by introducing the capability to generate and utilize probability distributions alongside proposed draft tokens. This allows for more sophisticated and accurate sampling during the drafting phase, improving the overall efficiency and quality of the speculative decoding process. Additionally, new tests have been introduced to ensure the stability of speculative decoding under complex scheduling conditions.

Highlights

• Speculative Decoding with Probabilities: The core change enables the Eagle speculative decoding method to not only propose draft tokens but also return their associated probability distributions. This moves beyond simple greedy (argmax) token generation to allow for more nuanced sampling strategies during drafting.
• Integration of Draft Probabilities: The GPUModelRunner now stores these draft probabilities per request in the CachedRequestState and passes them to the rejection sampler. This is crucial for the main model to efficiently verify and accept/reject the drafted tokens based on their likelihoods.
• Enhanced Testing for Speculative Decoding: New comprehensive tests have been added to ensure the robustness of speculative decoding under various scheduling scenarios, including partial scheduling and request preemption. This validates the stability of the new probability-aware drafting mechanism.
• Refactored Draft Proposal Interface: The propose_draft_token_ids method in GPUModelRunner has been renamed to propose_draft and its signature updated to explicitly return both the drafted token IDs and their probabilities, standardizing the interface for different speculative decoding methods.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request adds support for returning probabilities along with tokens in speculative decoding, which is a key requirement for rejection sampling. The changes are well-implemented across the gpu_model_runner and eagle proposer, with corresponding updates to tests. A new integration test for scheduling with speculative decoding is also a great addition. I have one minor suggestion to improve comment clarity for future maintainability.

[andylolu2]

@WoosukKwon PR is ready for review

[ekagra-ranjan]

Hi @andylolu2 - thanks for this PR! The AL improvement looks good.

One reason for using the argmax was to reduce the overhead of sampling from draft. TRTLLM and TGI also uses argmax. Could you also run the online benchmark on MTBench to see how the e2e gains in TPOT metric due to EAGLE look like with and without this PR? Example cmd: #18847

[andylolu2]

Hi @andylolu2 - thanks for this PR! The AL improvement looks good.

One reason for using the argmax was to reduce the overhead of sampling from draft. TRTLLM and TGI also uses argmax. Could you also run the online benchmark on MTBench to see how the e2e gains in TPOT metric due to EAGLE look like with and without this PR? Example cmd: #18847

@ekagra-ranjan I've added online benchmark numbers. The difference is not massive (results a bit noisy in general) at low temperature. I think it will still be good to merge this since:

1. Sampling can be improved quite easily (e.g. w/ a triton kernel)
2. The overall profile is more robust to changes in temperature.

[ekagra-ranjan]

@andylolu2 - thanks for adding the plot with TPOT. Can you also share the absolute numbers and the relative gain/degradation for different Temp?

[andylolu2]

@andylolu2 - thanks for adding the plot with TPOT. Can you also share the absolute numbers and the relative gain/degradation for different Temp?

The absolute numbers are in the "Raw numbers" collapse under the plot.

[ekagra-ranjan]

The drop in gains on low temp like 0.3 is ~10% which is a lot. Cohere for e.g., uses 0.3 as the default temperature. Coding and reasoning tasks usually use lower temperature which becomes even more important with thinking/reasoning models.

Can we have add an if-else to use the argmax method when the engine is using temp is < 0.75 to preserve the perf for these scenarios?

[andylolu2]

The drop in gains on low temp like 0.3 is ~10% which is a lot. Cohere for e.g., uses 0.3 as the default temperature. Coding and reasoning tasks usually use lower temperature which becomes even more important with thinking/reasoning models.

Can we have add an if-else to use the argmax method when the engine is using temp is < 0.75 to preserve the perf for these scenarios?

I think in general it's highly model-specific choice of what sampling temperature you should use for the draft model.
Sometimes you want to match the temperature of the target, sometimes you want to use higher/lower.

I suggest we make the threshold T configurable with the following heuristic:

1. When target temperature < T, sample draft model with temperature=0 (argmax).
2. When target tempearture >= T, sample draft model with temperature=T.

Does that sound reasonable to you?

[ekagra-ranjan]

Yes, we can have the threshold T as a parameter. Perhaps the default value should be 0.75 based on your results instead of having it a required param.

I think in general it's highly model-specific choice of what sampling temperature you should use for the draft model.
Sometimes you want to match the temperature of the target, sometimes you want to use higher/lower.

Oh, maybe I missed it but in your experiment are you using different temp for draft from target?

[andylolu2]

Oh, maybe I missed it but in your experiment are you using different temp for draft from target?

I'm using the same temperature for both.

[andylolu2]

@ekagra-ranjan I realised it's actually quite difficult to make my proposal work. Problem is the rejection sampler does not allow both:

1. Partially "unset draft probs" (i.e. some draft probs are filler values); and
2. Greedy draft sampling when target sampling is not greedy.

To make this work would require some large-ish amount of change to the rejection sampler, and that would be a rabbit hole to make sure I don't introduce unwanted overheads.

Instead I've optimized the draft model sampler a bit, the overhead is in the worst case ~4% now. New numbers updated in the PR description.

[ekagra-ranjan]

Problem is the rejection sampler does not allow both:

Do you think its simpler to select the old argmax approach only if all of the req in batch have temp below T OR instead select the new approach if all req in batch have temp above T?

Instead I've optimized the draft model sampler a bit, the overhead is in the worst case ~4% now. New numbers updated in the PR description.

Nice, Could you pls share which line of code/commit does this?

[andylolu2]

Do you think its simpler to select the old argmax approach only if all of the req in batch have temp below T OR instead select the new approach if all req in batch have temp above T?

I don't see an easy way. Due to the fact that drafts are not used immediately (usually the step right after, but can be in theory arbitrarily later on due to preemption), even if we do the fallback argmax approach during drafting, we might still end up with some requests with and some requests without draft probs during verification. A change to the rejection sampler would be needed.

Nice, Could you pls share which line of code/commit does this?

Forgot to push hahaa. It's this commit 20e43fd.

Also would like to add that the sampling overhead will be very negligible for larger models (e.g. DeepSeekV3/R1), so they should benefit a lot more from the increase in AL.

[keyboardAnt]

@jmamou, how does rebasing your eagle-frspec changes onto this branch affect eagle-frspec’s acceptance?

[ekagra-ranjan]

@andylolu2 - I was comparing the old and the new bench numbers

1. AL went from 2.29 to 2.18 at T=0. For reference: K=3, mtbench llama 3.18b gets 2.29
2. the overhead of w/ prob is ~9% for low temp like T=0.3 which is higher than the previously ~5%

Can you share why the new numbers look worse than the older ones?

[andylolu2]

@ekagra-ranjan I think the data I used wget https://raw.githubusercontent.com/SafeAILab/EAGLE/49b67c8d2dc6d8154eab6ef6899e4b9cc4cd3e06/eagle/data/mt_bench/question.jsonl -o mt_bench.jsonl is actually different from --dataset-name hf --dataset-path philschmid/mt-bench.

And for the same reason I switched the online benchmark numbers to use https://raw.githubusercontent.com/SafeAILab/EAGLE/49b67c8d2dc6d8154eab6ef6899e4b9cc4cd3e06/eagle/data/mt_bench/question.jsonl for easier comparison against the AL numbers.

I've also implemented the "kill switch" mechanism enable_draft_probs in SpeculativeConfig so if users see a regression they can always go back to greedy sampling.

[keyboardAnt]

@ekagra-ranjan I think the data I used wget https://raw.githubusercontent.com/SafeAILab/EAGLE/49b67c8d2dc6d8154eab6ef6899e4b9cc4cd3e06/eagle/data/mt_bench/question.jsonl -o mt_bench.jsonl is actually different from --dataset-name hf --dataset-path philschmid/mt-bench.

And for the same reason I switched the online benchmark numbers to use https://raw.githubusercontent.com/SafeAILab/EAGLE/49b67c8d2dc6d8154eab6ef6899e4b9cc4cd3e06/eagle/data/mt_bench/question.jsonl for easier comparison against the AL numbers.

@andylolu2, I believe the acceptance should be monotonically increasing regardless of the data (for exact expected acceptance rates, see sgl-project/sglang#9877).

Edit: Non-greedy sampling should have monotonically increasing acceptance over the acceptance of greedy sampling.

[ekagra-ranjan]

@andylolu2

Thanks for adding enable_draft_probs. This will be super useful for cases when sampling overhead eats into the gains!

I think its better to use --dataset-name hf --dataset-path philschmid/mt-bench since all previous benchmarks have been using it. The goal is to measure regression on vllm system so --dataset-name hf --dataset-path philschmid/mt-bench is the standard comparison.

[andylolu2]

Edit: Non-greedy sampling should have monotonically increasing acceptance over the acceptance of greedy sampling.

Yes agree, my numbers agree with that too.

[andylolu2]

@andylolu2

Thanks for adding enable_draft_probs. This will be super useful for cases when sampling overhead eats into the gains!

I think its better to use --dataset-name hf --dataset-path philschmid/mt-bench since all previous benchmarks have been using it. The goal is to measure regression on vllm system so --dataset-name hf --dataset-path philschmid/mt-bench is the standard comparison.

Makes sense, let me re-run the numbers for easier comparison across PRs.

[andylolu2]

@andylolu2

Thanks for adding enable_draft_probs. This will be super useful for cases when sampling overhead eats into the gains!

I think its better to use --dataset-name hf --dataset-path philschmid/mt-bench since all previous benchmarks have been using it. The goal is to measure regression on vllm system so --dataset-name hf --dataset-path philschmid/mt-bench is the standard comparison.

I've reran the benchmarks with --dataset-name hf --dataset-path philschmid/mt-bench

[LiuXiaoxuanPKU]

Thanks for the PR!! It's a great feature and promising numbers! I just have a concern about the seeding here, see detailed comment below.

[LiuXiaoxuanPKU]

I have a concern about this line:
The line q.exponential_() fills the tensor q with samples from an exponential distribution using PyTorch's global random number generator. This affects the global random seed and will advance the state of PyTorch's RNG. If the main model or other parts of code rely on reproducibility through seeding (e.g., torch.manual_seed(seed)), calling q.exponential_() here will consume random numbers from the global RNG, which may change subsequent random samples elsewhere.

[andylolu2]

@LiuXiaoxuanPKU do you have a proposed fix?
And I'm also not exactly sure in what circumstances it can break

[ekagra-ranjan]

And I'm also not exactly sure in what circumstances it can break

it will break reproducibility. vLLM users provide seed to the llm engine so that they can obtain the same result during reruns. This helps in investigating issues and putting llm in CI testing.

[luyuzhe111]

It looks like similar code exists already https://github.com/vllm-project/vllm/blob/main/vllm/v1/sample/rejection_sampler.py#L591. Does this mean this line is in rejection_sampler.py is problematic too?

[LiuXiaoxuanPKU]

For ngram, we can keep passing draft_probs as None to reduce the overhead?

[andylolu2]

This is inside _dummy_sampler_run, do you mean the warm up path might be different?

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @andylolu2.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[andylolu2]

@ekagra-ranjan @LiuXiaoxuanPKU Any updates? Can we make this PR merge-ready?

[jmamou]

@jmamou, how does rebasing your eagle-frspec changes onto this branch affect eagle-frspec’s acceptance?

@keyboardAnt
Acceptance rate reported at #24506

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @andylolu2.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[MLKoz2]

It is a nice change, would be nice to make this feature working.

[mergify]

Documentation preview: https://vllm--20459.org.readthedocs.build/en/20459/