diff --git a/vllm/model_executor/models/glm_ocr_mtp.py b/vllm/model_executor/models/glm_ocr_mtp.py
index 6c8e208c4b..05c3ac7935 100644
--- a/vllm/model_executor/models/glm_ocr_mtp.py
+++ b/vllm/model_executor/models/glm_ocr_mtp.py
@@ -86,7 +86,8 @@ class GlmOcrMultiTokenPredictorLayer(nn.Module):
     ) -> torch.Tensor:
         assert inputs_embeds is not None
         # masking inputs at position 0, as not needed by MTP
-        inputs_embeds[positions[0] == 0] = 0
+        token_positions = positions[0] if positions.ndim == 2 else positions
+        inputs_embeds.masked_fill_((token_positions == 0).unsqueeze(-1), 0)
 
         inputs_embeds = self.enorm(inputs_embeds)
         previous_hidden_states = self.hnorm(previous_hidden_states)