[Kernel] Remove hard-dependencies of Speculative decode to CUDA workers #10587

xuechendi · 2024-11-23T02:19:12Z

This PR is mainly target to remove hard dependency for CUDA in speculative decoding

Initiated this proposal in #10131

Original PR => use current_platform to provide Dynamic WorkerCls as BaseClass for spec_decode workers
This PR provides second solution => Use WorkerWrapperBase and ModerRunnerWrapperBase to lazily create spec decoder workers.

Test scripts:

pytest -v tests/spec_decode/e2e/test_mlp_correctness.py::test_mlp_e2e_greedy_correctness 
pytest -v tests/spec_decode/e2e/test_medusa_correctness.py::test_medusa_e2e_greedy_correctness
pytest -v tests/spec_decode/e2e/test_ngram_correctness.py::test_ngram_e2e_greedy_correctness
pytest -v tests/spec_decode/e2e/test_eagle_correctness.py::test_eagle_e2e_greedy_correctness

mergify · 2024-11-23T02:19:48Z

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

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

github-actions · 2024-11-23T02:20:13Z

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

🚀

xuechendi · 2024-11-23T02:27:59Z

@comaniac @youkaichao , based on #10555 , I submitted an alternative impl to derive from WorkerWrapperBase and ModelRunnerWrapperBase.

Please check if this makes sense to you.

youkaichao · 2024-11-24T07:30:27Z

Looks much better than previous one.

I'll hand it over to @comaniac

comaniac

Overall LGTM

comaniac · 2024-11-24T23:07:32Z

vllm/spec_decode/spec_decode_worker.py

+                    if current_platform.is_cuda_alike():
+                        draft_worker_kwargs[
+                            "model_runner_cls"] = TP1DraftModelRunner


So in non-CUDA platforms this key will be missing. Is this expected?

@comaniac , yes, TP1DraftModelRunner is a wrapper on ModelRunner, this wrapper will add extra data processing for
previous_hidden_states and other inputs before calling actual model_runner.execute_model()

In CPU enabling, we don't need extra work on previous_hidden_states to make it work, so CPUTP1DraftModelRunner is not necessary. However, in order to improve Spec Decode performance on other HWs, adding HWTP1DraftModelRunner might still be useful.

xuechendi · 2024-11-25T18:57:29Z

@comaniac , I have verified this PR on GPU and CPU in my local env.

youkaichao · 2024-11-25T21:44:47Z

vllm/platforms/cpu.py

+            if vllm_config.speculative_config:
+                parallel_config.worker_cls = \
+                    "vllm.spec_decode.spec_decode_worker.create_spec_worker"
+                parallel_config.actual_worker_cls = \


can we rename the field to sd_worker_cls ?

maybe add a field sd_worker_cls to ParallelConfig , and default to None.

youkaichao

LGTM now, thanks for the hard work!

mergify · 2024-11-26T17:13:50Z

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

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

mgoin · 2024-11-26T22:35:26Z

vllm/spec_decode/multi_step_worker.py

+if current_platform.is_cuda_alike():
+    from vllm.spec_decode.draft_model_runner import TP1DraftModelRunner


Do we have to put this check in the module import? It would be better if this was only lazy imported within sampler_output

mgoin · 2024-11-26T22:36:00Z

vllm/spec_decode/spec_decode_worker.py

+if current_platform.is_cuda_alike():
+    from vllm.spec_decode.draft_model_runner import TP1DraftModelRunner


Ditto but within create_worker

youkaichao · 2024-11-27T01:57:42Z

@xuechendi please address @mgoin 's comments, maybe in a followup PR.

xuechendi · 2024-11-27T01:58:37Z

@xuechendi please address @mgoin 's comments, maybe in a followup PR.

Got it, will do

Signed-off-by: Chendi Xue <[email protected]>

…project#10587) Signed-off-by: Chendi Xue <[email protected]> Signed-off-by: Andrew Feldman <[email protected]>

…project#10587) Signed-off-by: Chendi Xue <[email protected]>

xuechendi requested review from zhuohan123, youkaichao, alexm-redhat, comaniac and njhill as code owners November 23, 2024 02:19

mergify bot added the needs-rebase label Nov 23, 2024

xuechendi mentioned this pull request Nov 23, 2024

[Kernel]Generalize Speculative decode from Cuda #10094

Closed

xuechendi force-pushed the spec_decode_generalize branch from 8579790 to 9b49363 Compare November 23, 2024 02:25

mergify bot removed the needs-rebase label Nov 23, 2024

comaniac reviewed Nov 24, 2024

View reviewed changes

comaniac added the ready ONLY add when PR is ready to merge/full CI is needed label Nov 25, 2024

comaniac approved these changes Nov 25, 2024

View reviewed changes

youkaichao reviewed Nov 25, 2024

View reviewed changes

youkaichao approved these changes Nov 25, 2024

View reviewed changes

mergify bot added the needs-rebase label Nov 26, 2024

xuechendi force-pushed the spec_decode_generalize branch from 51e26f2 to 933e5e0 Compare November 26, 2024 19:00

xuechendi requested a review from LiuXiaoxuanPKU as a code owner November 26, 2024 19:00

xuechendi force-pushed the spec_decode_generalize branch from 933e5e0 to f58440a Compare November 26, 2024 19:06

mergify bot removed the needs-rebase label Nov 26, 2024

mgoin reviewed Nov 26, 2024

View reviewed changes

youkaichao merged commit 0a71900 into vllm-project:main Nov 27, 2024
49 of 51 checks passed

Spec Decode - Remove hard-dependency on GPU

90339c8

Signed-off-by: Chendi Xue <[email protected]>

xuechendi added 4 commits November 26, 2024 23:13

Remove Dynamic WorkerCls, use WorkerWrapperBase instead

e7b0ccd

Signed-off-by: Chendi Xue <[email protected]>

Enable CPU spec decode with small fix

5a0186e

Signed-off-by: Chendi Xue <[email protected]>

rename to sd_worker_cls

959c7ec

Signed-off-by: Chendi Xue <[email protected]>

Fix CI

f58440a

Signed-off-by: Chendi Xue <[email protected]>

afeldman-nm pushed a commit to neuralmagic/vllm that referenced this pull request Dec 2, 2024

Remove hard-dependencies of Speculative decode to CUDA workers (vllm-…

dd8ea8b

…project#10587) Signed-off-by: Chendi Xue <[email protected]> Signed-off-by: Andrew Feldman <[email protected]>

sleepwalker2017 pushed a commit to sleepwalker2017/vllm that referenced this pull request Dec 13, 2024

Remove hard-dependencies of Speculative decode to CUDA workers (vllm-…

55f8c8b

…project#10587) Signed-off-by: Chendi Xue <[email protected]>

xuechendi deleted the spec_decode_generalize branch December 19, 2024 21:49

BKitor pushed a commit to BKitor/vllm that referenced this pull request Dec 30, 2024

Remove hard-dependencies of Speculative decode to CUDA workers (vllm-…

7e2e400

…project#10587) Signed-off-by: Chendi Xue <[email protected]>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Kernel] Remove hard-dependencies of Speculative decode to CUDA workers #10587

[Kernel] Remove hard-dependencies of Speculative decode to CUDA workers #10587

xuechendi commented Nov 23, 2024 •

edited by github-actions bot

Loading

mergify bot commented Nov 23, 2024

github-actions bot commented Nov 23, 2024

xuechendi commented Nov 23, 2024

youkaichao commented Nov 24, 2024

comaniac left a comment

comaniac Nov 24, 2024

xuechendi Nov 25, 2024

xuechendi commented Nov 25, 2024

youkaichao Nov 25, 2024

youkaichao Nov 25, 2024

xuechendi Nov 25, 2024

youkaichao left a comment

mergify bot commented Nov 26, 2024

mgoin Nov 26, 2024

mgoin Nov 26, 2024

youkaichao commented Nov 27, 2024

xuechendi commented Nov 27, 2024

		if current_platform.is_cuda_alike():
		from vllm.spec_decode.draft_model_runner import TP1DraftModelRunner

[Kernel] Remove hard-dependencies of Speculative decode to CUDA workers #10587

[Kernel] Remove hard-dependencies of Speculative decode to CUDA workers #10587

Conversation

xuechendi commented Nov 23, 2024 • edited by github-actions bot Loading

mergify bot commented Nov 23, 2024

github-actions bot commented Nov 23, 2024

xuechendi commented Nov 23, 2024

youkaichao commented Nov 24, 2024

comaniac left a comment

Choose a reason for hiding this comment

comaniac Nov 24, 2024

Choose a reason for hiding this comment

xuechendi Nov 25, 2024

Choose a reason for hiding this comment

xuechendi commented Nov 25, 2024

youkaichao Nov 25, 2024

Choose a reason for hiding this comment

youkaichao Nov 25, 2024

Choose a reason for hiding this comment

xuechendi Nov 25, 2024

Choose a reason for hiding this comment

youkaichao left a comment

Choose a reason for hiding this comment

mergify bot commented Nov 26, 2024

mgoin Nov 26, 2024

Choose a reason for hiding this comment

mgoin Nov 26, 2024

Choose a reason for hiding this comment

youkaichao commented Nov 27, 2024

xuechendi commented Nov 27, 2024

xuechendi commented Nov 23, 2024 •

edited by github-actions bot

Loading