From 89dbb2644552d6e473a7e97da0ce8f0aa8e32c9d Mon Sep 17 00:00:00 2001 From: Yuxuan Zhang Date: Fri, 11 Sep 2026 20:31:10 +0800 Subject: [PATCH] [Bugfix] Fix GLM-OCR MTP position masking during CUDA graph capture (#56447) Signed-off-by: zRzRzRzRzRzRzR --- vllm/model_executor/models/glm_ocr_mtp.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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)