[V1][Spec Decode] Ngram Spec Decode #12193

LiuXiaoxuanPKU · 2025-01-19T23:17:33Z

This PR tries to add ngram spec decode to V1. Design doc: here.
Major changes:

Since we only implement the ngram spec decode, we did not add another scheduler for running the drafting method. We always check if we need to do ngram lookup before calling the scheduler.
Add a new field _spec_token_ids in Request to track speculated tokens.
Changes to model_runner:
3.1 Change the _prepare_input to also return the logits of speculated tokens.
3.2 Change the _prepare_input to add speculated tokens as input tokens.
3.3 Change the execute_model to generate multiple tokens per call. Concretely, it will add more than one tokens to input_batch and req_state.
We only perform spec decode for requests in the running queue.
We only support greedy decoding for now.

What is missing

Change scheduling to only propose tokens for decoding requests.
Stop checking for spec decode, where mutiple tokens are generated in a single step.
For the ngram lookup logic, currently I just append dummy tokens directly instead of performing the lookup. We can move v0's lookup logic here.
Check the correctness of this PR with chunked prefill. <-- We only perform spec decode in the decoding phase.
More end to end tests & Style.

Tasks out of the scope of this PR

Optimize the performance of ngram lookup.
Support non-greedy decoding.
Add other spec decode methods.

[Update]
I will move the following two features into following PRs:

Guarantee the correctness of prefix caching + spec decode, because it will involve changing the behavior of kv cache manager @comaniac.
Change the scheduling policy to guarantee that at least one token is scheduled for each request. Separate this because it will touch the scheduling code and needs more careful thought/test.

Minor: There is a minimal example/test in tests/v1/e2e/test_basic_specdecode.py. You can check it for the current use and check correctness with pytest -s tests/v1/e2e/test_basic_specdecode.py.

github-actions · 2025-01-19T23:17:44Z

👋 Hi! Thank you for contributing to the vLLM project.
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 do one of these:

Add ready label to the PR
Enable auto-merge.

🚀

mergify · 2025-01-19T23:18:10Z

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

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

Signed-off-by: LiuXiaoxuanPKU <lilyliupku@gmail.com>

NickLucche · 2025-01-24T14:26:52Z

Surely late here, but why is a speculative decoding-aware scheduler needed? Wouldn't it be possible to just assume multi-token generation per-step as default?

comaniac · 2025-01-24T16:18:15Z

Surely late here, but why is a speculative decoding-aware scheduler needed? Wouldn't it be possible to just assume multi-token generation per-step as default?

Because the scheduler has to know how many kv-cache slots are needed for each request. We use lookahead slots in v0 that always allocates k lookahead slots for each request when spec decode is enabled. However, it's inefficient when we don't have k spec tokens for every request. This may happen in the following examples:

N-gram won't propose any tokens if failed to find a match.
The draft model generates EOS.
Insufficient kv-cache slots.
In dynamic speculative decoding, we control "k" based on the current traffic.

So in this design for v1, we first get the spec tokens, and let the target model scheduler allocate the exact number of slots accordingly.

vllm/v1/core/scheduler.py

comaniac · 2025-01-24T16:22:59Z

vllm/v1/request.py

+        self._spec_token_ids = []
+
+    @property
+    def spec_token_ids(self) -> List[int]:


This function should return ConstantList to be read-only. See output_token_ids and all_token_ids as references.

vllm/v1/request.py

vllm/v1/core/scheduler.py

comaniac · 2025-01-24T17:36:10Z

vllm/v1/core/kv_cache_manager.py

+        # When calculating new full blocks, we exclude speculative tokens.
+        # We only cache blocks where token_ids are valid. KV cache of
+        # speculative tokens will be valid once these tokens are accepted
+        # (tracked by num_computed_tokens).
+        num_cached_tokens = request.num_computed_tokens + num_tokens - len(
+            request.spec_token_ids)
+        num_full_blocks_after_append = num_cached_tokens // self.block_size


I just realized that the current _cache_full_blocks may not support caching the same block twice. It may have issues in an edge case (block size 4, k=3):

Step 1: [0, 1, 2, 3] + [4, S0, S1, S2] - The first block is already cached in the last step. - Assuming all spec tokens are rejected. Step 2: [0, 1, 2, 3] + [4, 5, S0, S1] + [S2] - `5` is the bonus token. - num_cached_tokens = 5 + 1 - 3 = 3 - num_full_blocks_after_append = 3 // 4 = 0 - So you attempt to cache the first block again.

#12415 should fix this.

comaniac · 2025-01-24T17:36:51Z

vllm/v1/sample/rejection_sampler.py

+logger = init_logger(__name__)
+
+
+class RejectionSampler(nn.Module):


Can we use FlashInfer kernel for reject sampling?

Will add a TODO, saying we can replace it with flashinfer in the future.

vllm/v1/spec_decode/ngram_proposer.py

vllm/v1/worker/gpu_input_batch.py

vllm/v1/worker/gpu_model_runner.py

NickLucche · 2025-01-24T18:49:01Z

I see, thanks a lot for elaborating @comaniac!

vllm/v1/engine/core.py

vllm/v1/spec_decode/ngram_proposer.py

LiuXiaoxuanPKU · 2025-01-25T23:49:14Z

vllm/v1/core/scheduler.py

@@ -621,6 +663,8 @@ class SchedulerOutput:
    num_scheduled_tokens: Dict[str, int]
    total_num_scheduled_tokens: int
    scheduled_encoder_inputs: Dict[str, List[int]]
+    use_spec_decode: bool
+    scheduled_spec_decode_tokens: Dict[str, List[int]]


Maybe a more general approach:

no specific to spec decode: add a function to scheduler so it can schedule n new tokens (with token ids)

specific to spec decode: rejection sampling, some tokens might be rejected, some ways of rewinding

sroy745

Thanks for the PR !

vllm/v1/spec_decode/ngram_proposer.py

vllm/v1/engine/core.py

vllm/v1/sample/rejection_sampler.py

sroy745 · 2025-01-25T04:02:48Z

vllm/v1/worker/gpu_input_batch.py

@@ -361,11 +363,15 @@ def make_sampling_metadata(
            # TODO - Replace this with incremental update to output token
            # statistics.
            output_token_ids.append(req_id_output_token_ids[req_id])
+            if rejection_sampling:
+                assert req_id_to_spec_token_ids is not None
+                spec_token_ids.append(req_id_to_spec_token_ids[req_id])


Is it possible that there are no speculations for a req_id? If so do we need to do req_id_to_spec_token_ids.get(req_id, []) ?

sroy745 · 2025-01-27T04:09:41Z

vllm/v1/core/scheduler.py

-                request.append_output_token_ids(token_id)
-                num_new_tokens = 1
+            if request.num_computed_tokens >= request.num_tokens:
+                request.clear_spec_tokens()


I am wondering if we should clear the spec tokens always? What happens if the target model does not accept any of the spec tokens? Is that case being handled here?

vllm/v1/core/scheduler.py

tests/v1/core/test_stop_checking.py

mergify · 2025-01-27T04:28:20Z

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

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

LiuXiaoxuanPKU · 2025-01-31T23:56:18Z

vllm/v1/sample/rejection_sampler.py

@@ -0,0 +1,47 @@
+import torch


Use torch tensor or numpy to implement the first version of rejection sampling

LiuXiaoxuanPKU added 7 commits January 18, 2025 17:24

skeleton

6071a4b

runnable but incorrect

be798e7

fix

f0976dd

pass for simple non spec case

6039933

pass args and minor variable name bug fix

03cd3dd

minor

26ba690

minimal example

ca5e0dd

mergify bot added the needs-rebase label Jan 19, 2025

LiuXiaoxuanPKU added 2 commits January 19, 2025 19:49

Merge branch 'main' of github.com:LiuXiaoxuanPKU/vllm into ngram

704634c

minor

62012d1

mergify bot removed the needs-rebase label Jan 20, 2025

LiuXiaoxuanPKU added 10 commits January 19, 2025 22:28

format

7bd3f27

basic test

b0c5d25

minor

d5ee081

minor

4e11585

stop checking

f915eda

test for stop checking

bd8ac07

style and disable scheduling chunked requests

008a41e

signed-off-by

784b24a

Signed-off-by: LiuXiaoxuanPKU <lilyliupku@gmail.com>

ngram proposer

f3f6ebc

style and minor output token fix

5e7306e

LiuXiaoxuanPKU marked this pull request as ready for review January 24, 2025 03:37

LiuXiaoxuanPKU requested review from WoosukKwon, robertgshaw2-redhat, njhill, ywang96, comaniac and alexm-redhat as code owners January 24, 2025 03:37

comaniac reviewed Jan 24, 2025

View reviewed changes

LiuXiaoxuanPKU commented Jan 25, 2025

View reviewed changes

vllm/v1/engine/core.py Outdated Show resolved Hide resolved

LiuXiaoxuanPKU commented Jan 25, 2025

View reviewed changes

vllm/v1/spec_decode/ngram_proposer.py Outdated Show resolved Hide resolved

shahedy2276541 approved these changes Jan 25, 2025

View reviewed changes

LiuXiaoxuanPKU commented Jan 25, 2025

View reviewed changes

sroy745 reviewed Jan 27, 2025

View reviewed changes

mergify bot added the needs-rebase label Jan 27, 2025

partial cleanup & update the kmp

a26df8d

LiuXiaoxuanPKU commented Jan 31, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[V1][Spec Decode] Ngram Spec Decode #12193

[V1][Spec Decode] Ngram Spec Decode #12193

LiuXiaoxuanPKU commented Jan 19, 2025 •

edited by github-actions bot

Loading

github-actions bot commented Jan 19, 2025

mergify bot commented Jan 19, 2025

NickLucche commented Jan 24, 2025

comaniac commented Jan 24, 2025 •

edited

Loading

comaniac Jan 24, 2025

comaniac Jan 24, 2025

comaniac Jan 24, 2025

comaniac Jan 24, 2025

LiuXiaoxuanPKU Jan 31, 2025

NickLucche commented Jan 24, 2025

LiuXiaoxuanPKU Jan 25, 2025

sroy745 left a comment

sroy745 Jan 25, 2025

sroy745 Jan 27, 2025

mergify bot commented Jan 27, 2025

LiuXiaoxuanPKU Jan 31, 2025

		logger = init_logger(__name__)


		class RejectionSampler(nn.Module):

[V1][Spec Decode] Ngram Spec Decode #12193

Are you sure you want to change the base?

[V1][Spec Decode] Ngram Spec Decode #12193

Conversation

LiuXiaoxuanPKU commented Jan 19, 2025 • edited by github-actions bot Loading

github-actions bot commented Jan 19, 2025

mergify bot commented Jan 19, 2025

NickLucche commented Jan 24, 2025

comaniac commented Jan 24, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

NickLucche commented Jan 24, 2025

Choose a reason for hiding this comment

sroy745 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mergify bot commented Jan 27, 2025

Choose a reason for hiding this comment

LiuXiaoxuanPKU commented Jan 19, 2025 •

edited by github-actions bot

Loading

comaniac commented Jan 24, 2025 •

edited

Loading