Skip to content

Commit 4db5765

Browse files
authored
[Feature]Adapt GLM-5.2 (vllm-project#11264)
### What this PR does / why we need it? This PR adds support for GLM-5.2 on Ascend. This pull request updates `model_returns_tuple` in `llm_base_proposer.py` to support DeepSeek-family MTP models (`DeepSeekMTPModel`). Since DeepSeek MTP recycles the post-final-norm hidden state, its forward pass returns a tuple `(logit_hidden, recycle_hidden)`, whereas other MTP families return a single tensor. vllm's PR: vllm-project/vllm#45895 ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? - vLLM version: v0.23.0 - vLLM main: vllm-project/vllm@dc68bd8 Signed-off-by: jiajinzhu2 <jiajinzhu@huawei.com>
1 parent 2f1c76a commit 4db5765

2 files changed

Lines changed: 14 additions & 2 deletions

File tree

vllm_ascend/patch/worker/patch_deepseek_mtp.py

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -51,8 +51,12 @@ def forward(
5151
hidden_states = self.eh_proj(torch.cat([inputs_embeds, previous_hidden_states], dim=-1))
5252

5353
hidden_states, residual = self.mtp_block(positions=positions, hidden_states=hidden_states, residual=None)
54-
hidden_states = residual + hidden_states
55-
return hidden_states
54+
hidden_states = residual + hidden_states # pre-final-norm (logits hidden)
55+
# Recycle the post-final-norm hidden into the next draft step.
56+
# compute_logits applies shared_head (== final norm) to the pre-norm
57+
# element, so logits and the recycle each get exactly one final-norm.
58+
# Matches SGLang's deepseek_nextn.
59+
return hidden_states, self.shared_head(hidden_states)
5660

5761

5862
class AscendDeepSeekMTP(DeepSeekMTP):

vllm_ascend/spec_decode/llm_base_proposer.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1557,6 +1557,14 @@ def set_inputs_first_pass(
15571557
return total_num_output_tokens, token_indices_to_sample, new_cad, None
15581558

15591559
def model_returns_tuple(self) -> bool:
1560+
if self.method == "mtp":
1561+
# DeepSeek-family MTP (deepseek_mtp.py) recycles the post-final-
1562+
# norm hidden, so its forward returns (logit_hidden,
1563+
# recycle_hidden). Other MTP families return a single tensor.
1564+
draft_model_config = getattr(self, "draft_model_config", None)
1565+
hf_config = getattr(draft_model_config, "hf_config", None)
1566+
architectures = getattr(hf_config, "architectures", []) or []
1567+
return "DeepSeekMTPModel" in architectures
15601568
return self.method not in ("mtp", "draft_model", "dflash")
15611569

15621570
def attn_update_stack_num_spec_norm(

0 commit comments

Comments
 (0)