[Model][Speculative Decoding] DeepSeek MTP spec decode #12755

luccafong · 2025-02-04T23:58:11Z

Implement DeepSeek MTP: #12181 to support DeepSeek MTP layers for next n prediction.

github-actions · 2025-02-04T23:58:22Z

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

🚀

comaniac

Otherwise LGTM. It's pretty clean so no concerns.

comaniac · 2025-02-05T00:49:05Z

tests/spec_decode/e2e/test_mtp_correctness.py

Do you know how long does it take to run all tests in this file?

vllm/spec_decode/draft_model_runner.py

vllm/spec_decode/spec_decode_worker.py

vllm/transformers_utils/configs/__init__.py

vllm/transformers_utils/configs/deepseek_mtp.py

comaniac · 2025-02-05T01:09:17Z

vllm/worker/worker.py

@@ -71,7 +71,7 @@ def __init__(
            or (speculative_config.draft_model_config.model ==
                model_config.model) \
            or (speculative_config.draft_model_config.hf_config.model_type
-                not in ["medusa", "mlp_speculator", "eagle"]) \
+                not in ["medusa", "mlp_speculator", "eagle", "deepseek_mtp"]) \


Suggested change

not in ["medusa", "mlp_speculator", "eagle", "deepseek_mtp"]) \

not in ("medusa", "mlp_speculator", "eagle", "deepseek_mtp")) \

LiuXiaoxuanPKU · 2025-02-05T02:01:21Z

vllm/model_executor/models/deepseek_mtp.py

+        inputs_embeds: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)


QQ: where did we truncate the input_ids?

for 1st stage: position 0 is masked for MTP, but it only applies to k=1, I need to change the mask to the [position<=k-1],
for 2+ stage, previous tokens in last stage is marked pre-computed, this is a bit complicated for k>1 on different layers, need to look into.

in short, the current change works for k=1 (which deepseek v3 model set), but need more changes for k>1

Neo9061

Any way to put a MD file instructing examples on how to use the MTP for SD?

Especially,

The num_nextn_predict_layers is 1, can we specify speculation length more than 1? and what are requirements on formatting the draft model artifacts?
Is this code compatible with multi-node inference? assume so as the draft is loaded in single GPU?

Neo9061 · 2025-02-05T03:10:47Z

tests/spec_decode/e2e/test_mtp_correctness.py

+
+# max. number of speculative tokens: this corresponds to
+# num_heads in the config.json of the speculator model.
+MAX_SPEC_TOKENS = 3


The num_nextn_predict_layers in DeepSeek V3 has only 1. Will that mean you will reuse the MTP head if I specify MAX_SEC_TOKENS more than 1?

this is a test file on dummy model. num_speculative_tokens should be <= num_nextn_predict_layers, the transformer blocks are different in different steps. I am adding some assertion for this case when user pass higher number.

Is there a way to just re-use the MTP to predict tokens whose k > 1? as essentially they are the same right?

You can print some warning that this is not expected though.

Neo9061 · 2025-02-05T04:01:46Z

vllm/model_executor/models/deepseek_mtp.py

+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        config = vllm_config.model_config.hf_config
+        self.mtp_start_layer_idx = config.num_hidden_layers


shouldn't the mtp_start_layer_idx be num_hidden_layers -1?

num_hidden_layers is 61 in DeepSeek config. The index of last layer is 60.

https://huggingface.co/deepseek-ai/DeepSeek-V3/raw/main/model.safetensors.index.json the last layer is 61 which is the mtp layer.

https://huggingface.co/deepseek-ai/DeepSeek-V3/raw/main/model.safetensors.index.json the last layer is 61 which is the mtp layer.

I see, thanks for clarifying!

Neo9061 · 2025-02-05T16:39:01Z

@luccafong Sorry have to ask those questions as I hope to use your implementation.

Have you tested it e2e with VLLM's multi-node distributed inference setting? asking as I can only deploy the model in multi-node settings.
If I want to re-use the MTP head to do speculation length k > 1, what is the hacking implementation you would recommend to just make it work? as k=1 is too limited in my application.

benchislett · 2025-02-05T18:30:43Z

@luccafong I have been working on a similar implementation locally, and have faced a few challenges that I'm not sure are addressed here. Have you validated the acceptance rate for k=1 for real weights?

I believe that the final RMSNorm in the DeepSeekV3 main model is not necessary for speculative decoding since the hnorm already normalizes the previous hidden weights received from the main model. It's unclear to me how it is classified in the DeepSeek-V3 technical report, but I think that the norm might be included in the output head and therefore not normalized as input to the MTP module. Anecdotally, I observe a small increase in acceptance rate with this change.

Also, I have noticed the acceptance rate becomes very low (<50%) when I enable the recently added MLA attention. Have you noticed this also? I am not sure what could cause this, maybe it is a bug fixed in recent commits to vLLM. I would like to know if this is an issue for your implementation.

simon-mo · 2025-02-05T18:56:35Z

vllm/transformers_utils/configs/deepseek_v3.py

Do we need a copy of the configuration? I think for local usage we don't need this file if you have --trust-remote-code.

Ah i see it is used for deepseek mtp config...hmmm

I realized we don't need it if we add tokenizer config in both models, so going to remove these

Signed-off-by: Lu Fang <[email protected]>

luccafong · 2025-02-06T04:54:01Z

@luccafong I have been working on a similar implementation locally, and have faced a few challenges that I'm not sure are addressed here. Have you validated the acceptance rate for k=1 for real weights?

I believe that the final RMSNorm in the DeepSeekV3 main model is not necessary for speculative decoding since the hnorm already normalizes the previous hidden weights received from the main model. It's unclear to me how it is classified in the DeepSeek-V3 technical report, but I think that the norm might be included in the output head and therefore not normalized as input to the MTP module. Anecdotally, I observe a small increase in acceptance rate with this change.

Also, I have noticed the acceptance rate becomes very low (<50%) when I enable the recently added MLA attention. Have you noticed this also? I am not sure what could cause this, maybe it is a bug fixed in recent commits to vLLM. I would like to know if this is an issue for your implementation.

The accept rate is around 56% during my testing, MLA attention could lead to different branch,
https://github.com/luccafong/vllm/blob/ds_mtp/vllm/spec_decode/multi_step_worker.py#L98 I fixed in a later commit.

regarding the norm, thanks for pointing out, let me try adjusting to see if there is an improvement.

luccafong · 2025-02-06T04:57:29Z

@luccafong Sorry have to ask those questions as I hope to use your implementation.

Have you tested it e2e with VLLM's multi-node distributed inference setting? asking as I can only deploy the model in multi-node settings.

If I want to re-use the MTP head to do speculation length k > 1, what is the hacking implementation you would recommend to just make it work? as k=1 is too limited in my application.

1.Not tested with multi node settings; 2. We can reuse if you do some model processing, e.g. duplicate the weights to different layers, the hacky changes will not be proper since for n predict layers >1, and we do a k > n predict layers, it is difficult to decide which layer to forward multiple times.
Note for now as commented in the other thread, some changes are needed for K>1, I am working in progress, let me update with you if it works.

mergify · 2025-02-06T07:17:03Z

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

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

Neo9061 · 2025-02-06T17:46:27Z

vllm/model_executor/models/deepseek_mtp.py

+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+        assert inputs_embeds is not None
+        inputs_embeds[positions <= spec_step_index] = 0


Please excuse my multiple questions.

inputs_embeds[positions <= spec_step_index] = 0 is for pre-filling stage for each MTP head correct? as during draft model (MTP head) decoding stage, the inputs_embeds is a single hidden vector.

That is what I saw in EAGLE workflow. It firstly enters the code from here with num_steps being 1 for prefilling (that is where the mask is effective). Then the num_steps becomes to be speculation length k and inputs_embed for each forward pass is a single embed vector.

But I didn't see your logic is modified in

vllm/vllm/spec_decode/draft_model_runner.py

Line 271 in 467a96a

for step in range(num_steps):

to introduce spec_step_index, where do you introduce it then?

Neo9061 · 2025-02-06T17:53:29Z

vllm/model_executor/models/deepseek_mtp.py

+            str(idx):
+            DeepSeekMultiTokenPredictorLayer(
+                config,
+                f"{prefix}.layers.{idx}",


For my case of re-using the single MTP for k > 1. I will specify k to be the same as number of MTP heads.

In that case, can I simply hack f"{prefix}.layers.{idx}" to be models.layers.61 and self.num_mtp_layers to be k?

Neo9061 · 2025-02-06T17:55:50Z

@luccafong Sorry have to ask those questions as I hope to use your implementation.

Have you tested it e2e with VLLM's multi-node distributed inference setting? asking as I can only deploy the model in multi-node settings.

If I want to re-use the MTP head to do speculation length k > 1, what is the hacking implementation you would recommend to just make it work? as k=1 is too limited in my application.

1.Not tested with multi node settings; 2. We can reuse if you do some model processing, e.g. duplicate the weights to different layers, the hacky changes will not be proper since for n predict layers >1, and we do a k > n predict layers, it is difficult to decide which layer to forward multiple times. Note for now as commented in the other thread, some changes are needed for K>1, I am working in progress, let me update with you if it works.

Thank you for your reply! Are you testing with single node of H200 then to load the target model?

I just found a bug to use VLLM class to load target model via distributed inference (official Ray setup) with speculative decoding arguments passed in. Without those arguments, the target model is loaded successfully.

#12841

benchislett · 2025-02-06T20:25:25Z

vllm/model_executor/models/deepseek_mtp.py

+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+        assert inputs_embeds is not None
+        inputs_embeds[positions <= spec_step_index] = 0


I have found that this line might interfere with CUDA graph recording. I am unsure why, but removing this line allowed the draft acceptance rate to go up to 85% in my testing, from <50% prior. This issue was not present for me when MLA was disabled, or when Enforce-Eager was enabled. I believe there might be some setup code in the MLA attention for cuda graphs causing issues with graph replay.

Regardless of the cause, I am curious if you can reproduce this issue. Do you see improved acceptance rates when this line is omitted?

Tried, but got 58%, not much improvement though. could you share your impl, there could be some difference

Hi @luccafong , see my code here:
#12915

I hope this can clear up any inconsistency. Please let me know if you cannot reproduce using my implementation, or if you identify any issues.

WhatGhost · 2025-02-08T03:28:06Z

@luccafong
Nice work,inspired me a lot!
I have two very small problems that have been bothering me.

I notice you use luccafong/deepseek_mtp_main_random model . I wonder is it only a small size deepseek-like model for testing?
How can i run a deepseek-v3 model with mtp？ Just set speculative_model to deepseek-v3 it own?

Thanks!

LiuXiaoxuanPKU · 2025-02-10T07:09:11Z

@luccafong Nice work,inspired me a lot! I have two very small problems that have been bothering me.

I notice you use luccafong/deepseek_mtp_main_random model . I wonder is it only a small size deepseek-like model for testing?

How can i run a deepseek-v3 model with mtp？ Just set speculative_model to deepseek-v3 it own?

Thanks!

Yes you can just plugin in deepseek model path as shown below, the following is my setup on 8xH200

llm = LLM(
    model="/path/to/DeepSeek-R1",
    tensor_parallel_size=8,
    speculative_draft_tensor_parallel_size=1,
    max_model_len=8192, # If you have enough memory with your hardware, you can ignore this
    num_speculative_tokens=1, # only 1 is supported for now
)

luccafong requested review from njhill, LiuXiaoxuanPKU, DarkLight1337, ywang96, zhuohan123, youkaichao, alexm-redhat and comaniac as code owners February 4, 2025 23:58

mergify bot added the speculative-decoding label Feb 4, 2025

luccafong force-pushed the ds_mtp branch from fd7ff28 to d8a95cf Compare February 5, 2025 00:06

LiuXiaoxuanPKU self-assigned this Feb 5, 2025

luccafong force-pushed the ds_mtp branch from d8a95cf to 9de0bdf Compare February 5, 2025 00:27

comaniac reviewed Feb 5, 2025

View reviewed changes

LiuXiaoxuanPKU reviewed Feb 5, 2025

View reviewed changes

lambert0312 mentioned this pull request Feb 5, 2025

[Feature] DeepSeek V3 optimization sgl-project/sglang#2591

Open

18 tasks

luccafong force-pushed the ds_mtp branch 3 times, most recently from d5a9924 to 793ef4f Compare February 5, 2025 02:30

Neo9061 reviewed Feb 5, 2025

View reviewed changes

simon-mo reviewed Feb 5, 2025

View reviewed changes

luccafong added 3 commits February 5, 2025 15:34

support DeepSeek MTP spec decode

30af532

Signed-off-by: Lu Fang <[email protected]>

cleanup and add comments

12a153e

Signed-off-by: Lu Fang <[email protected]>

fix multi step and support same model checkpoint for main and spec model

9afc75e

luccafong force-pushed the ds_mtp branch from 793ef4f to 9afc75e Compare February 5, 2025 23:47

allow mla for gpu multi step spec decode

ce36703

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

Neo9061 reviewed Feb 6, 2025

View reviewed changes

benchislett reviewed Feb 6, 2025

View reviewed changes

This was referenced Feb 7, 2025

Add EAGLE-style MTP module support for DeepSeek-R1 and example usage script CentML/vllm#2

Draft

[Model][Speculative Decoding] Add EAGLE-style MTP module reference code for DeepSeek-R1 #12915

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Model][Speculative Decoding] DeepSeek MTP spec decode #12755

[Model][Speculative Decoding] DeepSeek MTP spec decode #12755

luccafong commented Feb 4, 2025

github-actions bot commented Feb 4, 2025

comaniac left a comment

comaniac Feb 5, 2025

comaniac Feb 5, 2025

LiuXiaoxuanPKU Feb 5, 2025

luccafong Feb 5, 2025

Neo9061 left a comment •

edited

Loading

Neo9061 Feb 5, 2025

luccafong Feb 5, 2025

Neo9061 Feb 5, 2025

Neo9061 Feb 5, 2025

luccafong Feb 5, 2025

Neo9061 Feb 5, 2025

Neo9061 commented Feb 5, 2025

benchislett commented Feb 5, 2025

simon-mo Feb 5, 2025

simon-mo Feb 5, 2025

luccafong Feb 5, 2025 •

edited

Loading

luccafong commented Feb 6, 2025

luccafong commented Feb 6, 2025 •

edited

Loading

mergify bot commented Feb 6, 2025

Neo9061 Feb 6, 2025

Neo9061 Feb 6, 2025

Neo9061 commented Feb 6, 2025

benchislett Feb 6, 2025

luccafong Feb 7, 2025

benchislett Feb 7, 2025

WhatGhost commented Feb 8, 2025

LiuXiaoxuanPKU commented Feb 10, 2025

	not in ["medusa", "mlp_speculator", "eagle", "deepseek_mtp"]) \
	not in ("medusa", "mlp_speculator", "eagle", "deepseek_mtp")) \

[Model][Speculative Decoding] DeepSeek MTP spec decode #12755

Are you sure you want to change the base?

[Model][Speculative Decoding] DeepSeek MTP spec decode #12755

Conversation

luccafong commented Feb 4, 2025

github-actions bot commented Feb 4, 2025

comaniac left a 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

Choose a reason for hiding this comment

Neo9061 left a comment • 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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Neo9061 commented Feb 5, 2025

benchislett commented Feb 5, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

luccafong Feb 5, 2025 • edited Loading

Choose a reason for hiding this comment

luccafong commented Feb 6, 2025

luccafong commented Feb 6, 2025 • edited Loading

mergify bot commented Feb 6, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Neo9061 commented Feb 6, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

WhatGhost commented Feb 8, 2025

LiuXiaoxuanPKU commented Feb 10, 2025

Neo9061 left a comment •

edited

Loading

luccafong Feb 5, 2025 •

edited

Loading

luccafong commented Feb 6, 2025 •

edited

Loading