[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.

Followup work:

benchmark the flashinfer rejection sampler
ngram kernel
proposer stop checking
KV cache based draft model

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 <[email protected]>

njhill

Thanks @LiuXiaoxuanPKU I think it's looking a lot better!

vllm/v1/request.py

vllm/v1/worker/gpu_model_runner.py

WoosukKwon

@LiuXiaoxuanPKU Thanks for the work! I think we are almost there. Left some comments, mostly regarding the code style.

vllm/v1/core/scheduler.py

vllm/v1/core/kv_cache_manager.py

vllm/v1/core/scheduler.py

vllm/v1/worker/gpu_model_runner.py

WoosukKwon · 2025-02-14T06:32:56Z

vllm/v1/worker/gpu_model_runner.py

+            if seq_len >= req_state.num_tokens:
+                # We don't rewind the generator state for requests now
+                # because spec decode only supports greedy decoding for now.


Could you please explain more?

It's not related, deleted.
But here, my main concern is how does spec decode handle seeding, rejected tokens might affect the generator state? Need more time to think about it and refer to V0.

Thanks for the fix!

WoosukKwon · 2025-02-15T07:50:02Z

@LiuXiaoxuanPKU Could you take a look at the failed tests?

I'll approve the PR once the tests are green!

mergify · 2025-02-15T11:59:52Z

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

WoosukKwon

@LiuXiaoxuanPKU Thanks for the great work! It's been quite a long journey. Really appreciate for all the work!

WoosukKwon · 2025-02-16T00:14:32Z

@LiuXiaoxuanPKU Just wanted to double check. Is the PR ready for merge?

JaheimLee · 2025-02-16T14:51:16Z

@WoosukKwon Hi, vllm nightly wheel doesn't have v1/spec_decode directory.

raceback (most recent call last):
  File "/data/lijinghui/uv_projects/.venv/lib/python3.12/site-packages/gunicorn/arbiter.py", line 608, in spawn_worker
    worker.init_process()
  File "/data/lijinghui/uv_projects/.venv/lib/python3.12/site-packages/uvicorn/workers.py", line 75, in init_process
    super().init_process()
  File "/data/lijinghui/uv_projects/.venv/lib/python3.12/site-packages/gunicorn/workers/base.py", line 135, in init_process
    self.load_wsgi()
  File "/data/lijinghui/uv_projects/.venv/lib/python3.12/site-packages/gunicorn/workers/base.py", line 147, in load_wsgi
    self.wsgi = self.app.wsgi()
                ^^^^^^^^^^^^^^^
  File "/data/lijinghui/uv_projects/.venv/lib/python3.12/site-packages/gunicorn/app/base.py", line 66, in wsgi
    self.callable = self.load()
                    ^^^^^^^^^^^
  File "/data/lijinghui/uv_projects/.venv/lib/python3.12/site-packages/gunicorn/app/wsgiapp.py", line 57, in load
    return self.load_wsgiapp()
           ^^^^^^^^^^^^^^^^^^^
  File "/data/lijinghui/uv_projects/.venv/lib/python3.12/site-packages/gunicorn/app/wsgiapp.py", line 47, in load_wsgiapp
    return util.import_app(self.app_uri)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/lijinghui/uv_projects/.venv/lib/python3.12/site-packages/gunicorn/util.py", line 370, in import_app
    mod = importlib.import_module(module)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mosh/.local/share/uv/python/cpython-3.12.7-linux-x86_64-gnu/lib/python3.12/importlib/__init__.py", line 90, in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<frozen importlib._bootstrap>", line 1387, in _gcd_import
  File "<frozen importlib._bootstrap>", line 1360, in _find_and_load
  File "<frozen importlib._bootstrap>", line 1331, in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 935, in _load_unlocked
  File "<frozen importlib._bootstrap_external>", line 995, in exec_module
  File "<frozen importlib._bootstrap>", line 488, in _call_with_frames_removed
  File "/data/lijinghui/uv_projects/LLM/chat_xiaoai.py", line 17, in <module>
    from vllm import SamplingParams, AsyncEngineArgs
  File "/data/lijinghui/uv_projects/.venv/lib/python3.12/site-packages/vllm/__init__.py", line 12, in <module>
    from vllm.engine.async_llm_engine import AsyncLLMEngine
  File "/data/lijinghui/uv_projects/.venv/lib/python3.12/site-packages/vllm/engine/async_llm_engine.py", line 1196, in <module>
    from vllm.v1.engine.async_llm import AsyncLLM
  File "/data/lijinghui/uv_projects/.venv/lib/python3.12/site-packages/vllm/v1/engine/async_llm.py", line 25, in <module>
    from vllm.v1.engine.core_client import EngineCoreClient
  File "/data/lijinghui/uv_projects/.venv/lib/python3.12/site-packages/vllm/v1/engine/core_client.py", line 20, in <module>
    from vllm.v1.engine.core import EngineCore, EngineCoreProc
  File "/data/lijinghui/uv_projects/.venv/lib/python3.12/site-packages/vllm/v1/engine/core.py", line 30, in <module>
    from vllm.v1.spec_decode.ngram_proposer import NgramProposer
ModuleNotFoundError: No module named 'vllm.v1.spec_decode'

WoosukKwon · 2025-02-16T16:02:05Z

@JaheimLee Thanks for reporting it! Fixed by #13359

Signed-off-by: LiuXiaoxuanPKU <[email protected]>

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 <[email protected]>

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

njhill reviewed Feb 13, 2025

View reviewed changes

vllm/v1/request.py Outdated Show resolved Hide resolved

vllm/v1/worker/gpu_model_runner.py Outdated Show resolved Hide resolved

vllm/v1/worker/gpu_model_runner.py Outdated Show resolved Hide resolved

vllm/v1/worker/gpu_model_runner.py Outdated Show resolved Hide resolved

markmc mentioned this pull request Feb 13, 2025

[Feature][v1]: Add metrics support #10582

Open

1 task

WoosukKwon reviewed Feb 14, 2025

View reviewed changes

LiuXiaoxuanPKU added 5 commits February 15, 2025 01:43

fix comments

8929ad1

minor fix

65bb67f

make test more deterministic

992aab8

Merge branch 'main' into ngram

6608e31

merge conflict

e298bb3

WoosukKwon mentioned this pull request Feb 15, 2025

[V1] Optimize handling of sampling metadata and req_ids list #13244

Merged

WoosukKwon added the ready ONLY add when PR is ready to merge/full CI is needed label Feb 15, 2025

LiuXiaoxuanPKU added 2 commits February 14, 2025 21:15

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

4329970

fix

4e015ae

mergify bot added the needs-rebase label Feb 15, 2025

LiuXiaoxuanPKU added 3 commits February 15, 2025 13:13

fix rejection sampler tests

5fc5264

fix num_token

b56a8e4

merge

a669c1c

mergify bot removed the needs-rebase label Feb 15, 2025

LiuXiaoxuanPKU added 3 commits February 15, 2025 13:55

fix scheduler test

2cbf57e

fix scheduler test, minor

29d3054

fix gpu model runner

2dc7909

WoosukKwon approved these changes Feb 16, 2025

View reviewed changes

WoosukKwon merged commit 80f63a3 into vllm-project:main Feb 16, 2025
40 of 42 checks passed

WoosukKwon mentioned this pull request Feb 16, 2025

[V1] Get input tokens from scheduler #13339

Merged

WoosukKwon mentioned this pull request Feb 16, 2025

[V1][BugFix] Add __init__.py to v1/spec_decode/ #13359

Merged

panf2333 pushed a commit to yottalabsai/vllm that referenced this pull request Feb 18, 2025

[V1][Spec Decode] Ngram Spec Decode (vllm-project#12193)

27b98bc

Signed-off-by: LiuXiaoxuanPKU <[email protected]>

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

njhill left a comment

WoosukKwon left a comment

WoosukKwon Feb 14, 2025

LiuXiaoxuanPKU Feb 15, 2025 •

edited

Loading

WoosukKwon Feb 15, 2025

WoosukKwon commented Feb 15, 2025

mergify bot commented Feb 15, 2025

WoosukKwon left a comment

WoosukKwon commented Feb 16, 2025

JaheimLee commented Feb 16, 2025 •

edited

Loading

WoosukKwon commented Feb 16, 2025

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

[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

njhill left a comment

Choose a reason for hiding this comment

WoosukKwon left a comment

Choose a reason for hiding this comment

WoosukKwon Feb 14, 2025

Choose a reason for hiding this comment

LiuXiaoxuanPKU Feb 15, 2025 • edited Loading

Choose a reason for hiding this comment

WoosukKwon Feb 15, 2025

Choose a reason for hiding this comment

WoosukKwon commented Feb 15, 2025

mergify bot commented Feb 15, 2025

WoosukKwon left a comment

Choose a reason for hiding this comment

WoosukKwon commented Feb 16, 2025

JaheimLee commented Feb 16, 2025 • edited Loading

WoosukKwon commented Feb 16, 2025

LiuXiaoxuanPKU commented Jan 19, 2025 •

edited by github-actions bot

Loading

LiuXiaoxuanPKU Feb 15, 2025 •

edited

Loading

JaheimLee commented Feb 16, 2025 •

edited

Loading