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 14, 00:30 UTC
[MTP] Refactor mtp predictor to avoid d2h operation (#27643)
Signed-off-by: MengqingCao <cmq0113@163.com>
Mengqing Cao committed
Oct 31, 2025 at 01:27 UTC
10042057953cd1528701234925de3d7b109e26de
1 file changed
+1
-1
vllm/model_executor/models/deepseek_mtp.py
+1
-1
@@ -97,7 +97,7 @@ class DeepSeekMultiTokenPredictorLayer(nn.Module):
97
) -> torch.Tensor:
98
assert inputs_embeds is not None
99
# masking inputs at position 0, as not needed by MTP
100
- inputs_embeds[positions == 0] = 0
100
+ inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds)
101
inputs_embeds = self.enorm(inputs_embeds)
102
previous_hidden_states = self.hnorm(previous_hidden_states)
103