[Bugfix] Fix GLM-OCR MTP position masking during CUDA graph capture (#56447)
Signed-off-by: zRzRzRzRzRzRzR <Yuxuan.Zhang2@liverpool.ac.uk>
Yuxuan Zhang committed
Sep 11, 2026 at 20:31 UTC
89dbb2644552d6e473a7e97da0ce8f0aa8e32c9d
1 file changed
+2
-1
vllm/model_executor/models/glm_ocr_mtp.py
+2
-1
@@ -86,7 +86,8 @@ class GlmOcrMultiTokenPredictorLayer(nn.Module):
86
) -> torch.Tensor:
87
assert inputs_embeds is not None
88
# masking inputs at position 0, as not needed by MTP
89
- inputs_embeds[positions[0] == 0] = 0
89
+ token_positions = positions[0] if positions.ndim == 2 else positions
90
+ inputs_embeds.masked_fill_((token_positions == 0).unsqueeze(-1), 0)
91
92
inputs_embeds = self.enorm(inputs_embeds)
93
previous_hidden_states = self.hnorm(previous_hidden_states)