Commit 755d78b4 authored by zhuwenwen's avatar zhuwenwen
Browse files

Merge branch 'v0.9.2-dev-wm-1127' into 'v0.9.2-dev'

[fix]解决宽松mtp1报错

See merge request dcutoolkit/deeplearing/vllm!272
parents afd18c9c deae0a22
......@@ -239,13 +239,18 @@ class EagleProposer:
if envs.VLLM_REJECT_SAMPLE_OPT:
draft_prob = logits.softmax(dim=-1, dtype=torch.float32)
draft_probs_list = [draft_prob]
# Early exit if there is only one draft token to be generated.
if self.num_speculative_tokens == 1:
# [batch_size, 1]
if envs.VLLM_REJECT_SAMPLE_OPT:
return draft_token_ids.view(-1, 1), draft_prob.view(-1, 1, logits.shape[-1])
return draft_token_ids.view(-1, 1)
if envs.VLLM_REJECT_SAMPLE_OPT:
draft_probs_list = [draft_prob]
# TODO: Currently, MTP module released by deepseek only has
# one layer. Adapt this code to support multiple layers once
# there's a multi-layer MTP module.
......
......@@ -166,13 +166,18 @@ class V1ZeroEagleProposer(EagleProposer):
if envs.VLLM_REJECT_SAMPLE_OPT:
draft_prob = logits.softmax(dim=-1, dtype=torch.float32)
draft_probs_list = [draft_prob]
# Early exit if there is only one draft token to be generated.
if self.num_speculative_tokens == 1:
# [batch_size, 1]
if envs.VLLM_REJECT_SAMPLE_OPT:
return draft_token_ids.view(-1, 1), draft_prob.view(-1, 1, logits.shape[-1])
return draft_token_ids.view(-1, 1)
if envs.VLLM_REJECT_SAMPLE_OPT:
draft_probs_list = [draft_prob]
# TODO: Currently, MTP module released by deepseek only has
# one layer. Adapt this code to support multiple layers once
# there's a multi-layer MTP module.
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment