A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 17:45 UTC
[Bugfix] Fix DeepSeek V4 MTP HC state handling (#42320)
Signed-off-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>
Mohammad Miadh Angkad committed
May 14, 2026 at 06:44 UTC
f1cc7aad3c2c09ff101715506ef36abedb7a0192
2 files changed
+8
-5
vllm/model_executor/models/deepseek_v4.py
+4
-4
@@ -1212,10 +1212,10 @@ class DeepseekV4DecoderLayer(nn.Module):
1212
x: torch.Tensor,
1213
positions: torch.Tensor,
1214
input_ids: torch.Tensor | None,
1215
- post_mix: torch.Tensor | None,
1216
- res_mix: torch.Tensor | None,
1217
- residual: torch.Tensor | None,
1218
- ) -> torch.Tensor:
1215
+ post_mix: torch.Tensor | None = None,
1216
+ res_mix: torch.Tensor | None = None,
1217
+ residual: torch.Tensor | None = None,
1218
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
1219
if residual is None:
1220
# Run standalone hc_pre on first layer
1221
residual = x
vllm/model_executor/models/deepseek_v4_mtp.py
+4
-1
@@ -143,9 +143,12 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
143
hidden_states = self.h_proj(previous_hidden_states) + self.e_proj(
144
inputs_embeds
145
).unsqueeze(-2)
146
- hidden_states = self.mtp_block(
146
+ hidden_states, residual, post_mix, res_mix = self.mtp_block(
147
positions=positions, x=hidden_states, input_ids=None
148
)
149
+ hidden_states = self.mtp_block.hc_post(
150
+ hidden_states, residual, post_mix, res_mix
151
+ )
152
# Return the flat pre-hc_head residual so it can be re-fed as the
153
# next spec step's `previous_hidden_states` when
154
# num_speculative_tokens > 1. hc_head is deferred to compute_logits.