[Speculative decoding 6/9] Integrate speculative decoding with LLMEngine #3894

cadedaniel · 2024-04-07T07:09:38Z

This PR integrates the existing speculative decoding components with the LLMEngine. With this PR, one can run speculative decoding with a draft model / target model e2e. Note however that the output is garbage as there remains e2e correctness implementation (PR 7/9 in the open sourcing plan).

Notes:

Testing: This PR relies on e2e tests with speculative decoding enabled. We compare that the number of emitted tokens is correct, and that the detokenized string matches what HF transformers produces.
Testing: this PR includes some unit tests for new detokenization/eos management in LLMEngine
LLMEngine Refactor: Because it is complicated to combine beam search logic with lookahead scheduling, we introduce a new interface inside the LLMEngine that allows us to process new token ids separately for beam search and lookahead scheduling.
Spec decode support: Currently the e2e test only passes for TP=1 on GPU. The work to add TP>1 support is trivial, but will be wasted as we need to modify the interfaces for TP>1 anyways. It is left out of this PR.
This PR modifies the WorkerBase interface slightly, such that all workers return List[SamplerOutput] from execute_model instead of just SamplerOutput.
This PR is stacked on top of [Misc] [Core] Implement RFC "Augment BaseExecutor interfaces to enable hardware-agnostic speculative decoding" #3837. A mirror PR has been created for easy viewing until that PR is merged.

LiuXiaoxuanPKU

Thank for the PR! Just some minor comments.

LiuXiaoxuanPKU · 2024-04-10T04:23:50Z

vllm/engine/output_processor/single_step.py

+                # Remove the parent sequence if it is not selected for next
+                # iteration
+                seq_group.remove(seq.seq_id)
+                self.scheduler.free_seq(seq)


Currently, output processor owns the scheduler and calls scheduler's functions in processing the output. This is a bit wired. Is it possible to move the scheduler part out of the output processor? Or is it too much work?

Agree! Note this weirdness is not introduced in this PR, it already existed in the LLMEngine. If it had unit tests we could easily refactor it, but I'm afraid without good unit tests it's risky to refactor things without a good reason.

LiuXiaoxuanPKU · 2024-04-10T04:30:59Z

vllm/engine/output_processor/stop_checker.py

+    """
+
+    def __init__(self, scheduler, scheduler_config, get_tokenizer_for_seq):
+        self.scheduler = scheduler


scheduler is not used in this class, why do we need scheduler here?

we need to read max_model_len from scheduler config, can we pass max_model_len directly?
In general, we might want to avoid the dependency on scheduler/scheduler_config in this class.

removed scheduler dependency and added type annotations (good catch!)

passed max model len directly

LiuXiaoxuanPKU · 2024-04-10T04:40:08Z

vllm/spec_decode/spec_decode_worker.py

        # If no spec tokens, call the proposer and scorer workers normally.
        # Used for prefill.
-        if num_spec_tokens == 0 or len(seq_group_metadata_list) == 0:
+        if num_lookahead_slots == 0 or len(seq_group_metadata_list) == 0:


When can len(seq_group_metadata_list) be 0?

I don't have any use in mind for it, but I want to maintain compatibility with the other workers.

cpu worker

vllm/vllm/worker/cpu_worker.py

Lines 284 to 286 in 92cd2e2

# If there is no input, we don't need to execute the model.

if num_seq_groups == 0:

return {}

gpu worker

vllm/vllm/worker/worker.py

Lines 231 to 233 in 92cd2e2

# If there is no input, we don't need to execute the model.

if num_seq_groups == 0:

return {}

neuron worker

vllm/vllm/worker/neuron_worker.py

Lines 79 to 81 in 92cd2e2

# If there is no input, we don't need to execute the model.

if num_seq_groups == 0:

return {}

LiuXiaoxuanPKU · 2024-04-10T21:03:51Z

vllm/engine/output_processor/multi_step.py

@@ -0,0 +1,117 @@
+from typing import Callable, Iterable, List


Another question:
From my understanding, multi_step processor is used for target model, which is not a muti_step_worker.
single_step processor is used for the draft model, which is a multi_step_worker.
Correct me if I'm wrong. If this is the case, can we have a better name for multi_step output processor?

The LLMEngine only interacts with the SpecDecodeWorker, which internally manages workers for the draft and target models. See this diagram from the meetup below:

This means that from the LLMEngine's perspective, there is only one worker type, and only one output processor type at a time.

cadedaniel · 2024-04-11T02:24:26Z

update: seems there is some conflict with #3884 and block manager v2 that only shows up in the spec decode e2e test added in this PR. will investigate.

rkooo567 · 2024-04-11T07:29:25Z

^ sorry for the inconvenience @cadedaniel!

cadedaniel · 2024-04-11T07:44:49Z

np glad chunked prefill got in

leiwen83 · 2024-04-11T08:30:44Z

Hi @cadedaniel ,

Great job for this feature enabling!
I locally pull your draft patch #3951, and test it with both drafter/verifer use the same model. The correctness of output is confirmed!

However, I find that in https://github.com/vllm-project/vllm/blob/main/vllm/spec_decode/multi_step_worker.py#L347, you append -1 towards the proposal to fill num_lookahead_slots number. Which maybe a bit of waste for the verification stage? Since the proposal may not always num_lookahead_slots length? Is there any improvement space here?

Thx~

cadedaniel · 2024-04-11T18:01:24Z

Hi @cadedaniel ,

Great job for this feature enabling! I locally pull your draft patch #3951, and test it with both drafter/verifer use the same model. The correctness of output is confirmed!

Thanks for testing it out.

However, I find that in https://github.com/vllm-project/vllm/blob/main/vllm/spec_decode/multi_step_worker.py#L347, you append -1 towards the proposal to fill num_lookahead_slots number. Which maybe a bit of waste for the verification stage? Since the proposal may not always num_lookahead_slots length? Is there any improvement space here?

Thx~

Yes, it would be wasteful if we verified these padding token ids. We avoid this via two means:

num_lookahead_slots is always k during speculation.
a proposal is either of len k or 0 -- we either speculate k tokens ahead of do not speculate at all.

This means we never have any -1 proposal tokens to verify.

In the future, we can support per-sequence proposal lens.

…ine (vllm-project#3894)

cadedaniel added 30 commits April 3, 2024 14:17

wip

252a0c7

Merge remote-tracking branch 'upstream/main' into executor_base

dd629d4

wip

a34800f

wip

09f30bd

clean

8b5bb8b

wip

6fd424f

wip

2a347bb

wip

658ff9b

wip

acee7be

wip

85760d6

wip

408b29d

Merge remote-tracking branch 'upstream/main' into executor_base

9d8fd69

wip

3149a03

wip

0c32e0a

wip

f64d5b1

wip

7207f0c

wip

0c4df0b

wip

2e355e7

wip

edb7f62

wip

48bb3e9

fix test

7b39044

fix test

9e5f2fb

fix test

1a3e26e

fix test

cd2015c

fix

d926034

fix

607f7e2

fix

e127bb7

fix

deaa8b0

clean

7817d61

clean

99823a3

cadedaniel mentioned this pull request Apr 9, 2024

[Speculative decoding 7/9] Speculative decoding end-to-end correctness tests. #3951

Merged

LiuXiaoxuanPKU reviewed Apr 10, 2024

View reviewed changes

cadedaniel added 2 commits April 10, 2024 14:13

pr feedback

de16919

Merge branch 'main' into llm-engine-spec

d933e50

cadedaniel enabled auto-merge (squash) April 11, 2024 01:35

cadedaniel added 6 commits April 16, 2024 02:37

allow append empty tokens in block table

2a19f5e

Merge remote-tracking branch 'upstream/main' into llm-engine-spec

79325d3

rebase on stop string fixes

b6e9e82

test spec

bf0c37c

lint & mypy

a158256

doc

5a69f6c

LiuXiaoxuanPKU approved these changes Apr 16, 2024

View reviewed changes

cadedaniel merged commit e95cd87 into vllm-project:main Apr 16, 2024
46 checks passed

cadedaniel mentioned this pull request Apr 16, 2024

[Misc] [CI] Fix CI failure caught after merge #4126

Merged

ronensc mentioned this pull request Apr 17, 2024

[Bugfix][Core] Restore logging of stats in the async engine #4150

Merged

robertgshaw2-neuralmagic pushed a commit to neuralmagic/nm-vllm that referenced this pull request Apr 21, 2024

[Speculative decoding 6/9] Integrate speculative decoding with LLMEng…

07a4d87

…ine (vllm-project#3894)

z103cb pushed a commit to z103cb/opendatahub_vllm that referenced this pull request Apr 22, 2024

[Speculative decoding 6/9] Integrate speculative decoding with LLMEng…

5ddd35c

…ine (vllm-project#3894)

robertgshaw2-neuralmagic pushed a commit to neuralmagic/nm-vllm that referenced this pull request Apr 26, 2024

[Speculative decoding 6/9] Integrate speculative decoding with LLMEng…

945a6b7

…ine (vllm-project#3894)

dtrifiro mentioned this pull request May 15, 2024

bump ubi base image tag opendatahub-io/vllm#24

Merged

Temirulan pushed a commit to Temirulan/vllm-whisper that referenced this pull request Sep 6, 2024

[Speculative decoding 6/9] Integrate speculative decoding with LLMEng…

e728793

…ine (vllm-project#3894)

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Speculative decoding 6/9] Integrate speculative decoding with LLMEngine #3894

[Speculative decoding 6/9] Integrate speculative decoding with LLMEngine #3894

cadedaniel commented Apr 7, 2024 •

edited

Loading

LiuXiaoxuanPKU left a comment

LiuXiaoxuanPKU Apr 10, 2024 •

edited

Loading

cadedaniel Apr 10, 2024

LiuXiaoxuanPKU Apr 10, 2024

cadedaniel Apr 10, 2024

LiuXiaoxuanPKU Apr 10, 2024

cadedaniel Apr 10, 2024

LiuXiaoxuanPKU Apr 10, 2024

cadedaniel Apr 10, 2024

cadedaniel commented Apr 11, 2024

rkooo567 commented Apr 11, 2024

cadedaniel commented Apr 11, 2024

leiwen83 commented Apr 11, 2024 •

edited

Loading

cadedaniel commented Apr 11, 2024

	# If there is no input, we don't need to execute the model.
	if num_seq_groups == 0:
	return {}

		@@ -0,0 +1,117 @@
		from typing import Callable, Iterable, List

[Speculative decoding 6/9] Integrate speculative decoding with LLMEngine #3894

[Speculative decoding 6/9] Integrate speculative decoding with LLMEngine #3894

Conversation

cadedaniel commented Apr 7, 2024 • edited Loading

LiuXiaoxuanPKU left a comment

Choose a reason for hiding this comment

LiuXiaoxuanPKU Apr 10, 2024 • edited Loading

Choose a reason for hiding this comment

cadedaniel Apr 10, 2024

Choose a reason for hiding this comment

LiuXiaoxuanPKU Apr 10, 2024

Choose a reason for hiding this comment

cadedaniel Apr 10, 2024

Choose a reason for hiding this comment

LiuXiaoxuanPKU Apr 10, 2024

Choose a reason for hiding this comment

cadedaniel Apr 10, 2024

Choose a reason for hiding this comment

LiuXiaoxuanPKU Apr 10, 2024

Choose a reason for hiding this comment

cadedaniel Apr 10, 2024

Choose a reason for hiding this comment

cadedaniel commented Apr 11, 2024

rkooo567 commented Apr 11, 2024

cadedaniel commented Apr 11, 2024

leiwen83 commented Apr 11, 2024 • edited Loading

cadedaniel commented Apr 11, 2024

cadedaniel commented Apr 7, 2024 •

edited

Loading

LiuXiaoxuanPKU Apr 10, 2024 •

edited

Loading

leiwen83 commented Apr 11, 2024 •

edited

Loading