Fix GLM-OCR MTP multimodal embeddings and positions (#39088)

Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com>
Co-authored-by: Xinyuan Tong <xinyuantong.cs@gmail.com>

python/sglang/srt/layers/rotary_embedding/mrope.py

@@ -250,6 +250,8 @@ class MRotaryEmbedding(RotaryEmbedding):
         fused_set_kv_buffer_arg=None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         assert positions.ndim == 1 or positions.ndim == 2
+        if query.shape[0] == 0:
+            return query, key
         self._match_cos_sin_cache_dtype(query)
         if positions.ndim == 2 and self.mrope_section:
             return self.forward_triton(positions, query, key)

python/sglang/srt/managers/mm_utils.py

@@ -746,7 +746,12 @@ def general_mm_embed_routine(
                                         )
                                     )
             forward_batch.mm_inputs = None
-            forward_batch.mm_input_embeds = input_embeds
+            forward_batch.mm_input_embeds = (
+                input_embeds.clone()
+                if forward_batch.spec_algorithm is not None
+                and forward_batch.spec_algorithm.is_eagle()
+                else input_embeds
+            )
         else:
             input_embeds = embed_tokens(input_ids)
         # Copy to pre-allocated buffer if available (for CUDA graph address stability)

python/sglang/srt/models/glm4v.py

@@ -667,7 +667,7 @@ class Glm4vForConditionalGeneration(nn.Module):
                 otherwise it will be `(seq_len,).
                 (Use input_metadata.mrope_positions to replace it)
         """
-        if self.is_mrope_enabled:
+        if self.is_mrope_enabled and forward_batch.mrope_positions is not None:
             positions = forward_batch.mrope_positions
 
         if not (

python/sglang/srt/models/glm_ocr_nextn.py

@@ -36,6 +36,7 @@ from sglang.srt.models.glm4 import Glm4DecoderLayer
 from sglang.srt.models.glm_ocr import GlmOcrForConditionalGeneration
 from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import add_prefix
+from sglang.srt.utils.hf_transformers_utils import get_rope_config
 
 logger = logging.getLogger(__name__)
 
@@ -86,9 +87,20 @@ class GlmOcrModelNextN(nn.Module):
         input_embeds: torch.Tensor = None,
     ) -> torch.Tensor:
         if input_embeds is None:
-            hidden_states = self.embed_tokens(input_ids)
-        else:
-            hidden_states = input_embeds
+            input_embeds = forward_batch.mm_input_embeds
+            if (
+                forward_batch.forward_mode.is_extend()
+                and forward_batch.contains_mm_inputs()
+                and not forward_batch.forward_mode.is_draft_extend_v2()
+            ):
+                assert input_embeds is not None
+                last_indices = (
+                    forward_batch.extend_start_loc + forward_batch.extend_seq_lens - 1
+                ).long()
+                input_embeds[last_indices] = self.embed_tokens(input_ids[last_indices])
+            if input_embeds is None:
+                input_embeds = self.embed_tokens(input_ids)
+        hidden_states = input_embeds
 
         if hidden_states.shape[0] > 0:
             hidden_states = self.eh_proj(
@@ -127,6 +139,8 @@ class GlmOcrForConditionalGenerationNextN(GlmOcrForConditionalGeneration):
         self.config = config
         self.tp_size = get_parallel().tp_size
         self.quant_config = quant_config
+        _, rope_scaling = get_rope_config(config)
+        self.is_mrope_enabled = "mrope_section" in (rope_scaling or {})
         self.model = GlmOcrModelNextN(
             config, quant_config, prefix=add_prefix("model", prefix)
         )
@@ -148,6 +162,8 @@ class GlmOcrForConditionalGenerationNextN(GlmOcrForConditionalGeneration):
         positions: torch.Tensor,
         forward_batch: ForwardBatch,
     ) -> torch.Tensor:
+        if self.is_mrope_enabled and forward_batch.mrope_positions is not None:
+            positions = forward_batch.mrope_positions
         hidden_states = self.model(input_ids, positions, forward_batch)
         return self.logits_processor(
             input_ids, hidden_states, self.lm_head, forward_batch

python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py

@@ -595,6 +595,11 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
             forward_batch.spec_info.topk_index,
             forward_batch.req_pool_indices,
         ]
+        if self.model_runner.model_config.model_is_mrope:
+            buffers.mrope_positions.zero_()
+            if forward_batch.mrope_positions is not None:
+                copy_dsts.append(buffers.mrope_positions[:, :raw_num_token])
+                copy_srcs.append(forward_batch.mrope_positions)
         if buffers.rids_int is not None and forward_batch.rids_int is not None:
             copy_dsts.append(buffers.rids_int[:raw_bs])
             copy_srcs.append(forward_batch.rids_int)

python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py

@@ -549,6 +549,11 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
             forward_batch.positions,
             forward_batch.req_pool_indices,
         ]
+        if self.model_runner.model_config.model_is_mrope:
+            buffers.mrope_positions.zero_()
+            if forward_batch.mrope_positions is not None:
+                copy_dsts.append(buffers.mrope_positions[:, :num_tokens])
+                copy_srcs.append(forward_batch.mrope_positions)
         if forward_batch.extend_seq_lens is not None:
             copy_dsts.append(buffers.extend_seq_lens[:raw_bs])
             copy_srcs.append(forward_batch.extend_seq_lens)

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -910,6 +910,8 @@ class EagleDraftWorker(EagleDraftWorkerBase):
                     )
                     topk_p, topk_index = fast_topk(probs, self.topk, dim=-1)
                     forward_batch.positions.add_(1)
+                if self.draft_runner.model_config.model_is_mrope:
+                    forward_batch.mrope_positions.add_(1)
                 maybe_detect_oob(
                     topk_index,
                     0,

python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py

@@ -274,6 +274,7 @@ class TinyModelConfig:
         self.swa_v_head_dim = head_dim
         self.is_encoder_decoder = False
         self.is_multimodal = False
+        self.model_is_mrope = False
         self.is_generation = True
         self.quantization = None
         self.is_hybrid_swa = sliding_window_size is not None

python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py

@@ -240,6 +240,7 @@ class TinyDSAModelConfig:
         self.kv_lora_rank = kv_lora_rank
         self.is_encoder_decoder = False
         self.is_multimodal = False
+        self.model_is_mrope = False
         self.is_generation = True
         self.quantization = None
         self.is_hybrid_swa = False

python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py

@@ -264,6 +264,7 @@ class TinyDSV4ModelConfig:
         self.sliding_window_size = DSV4_SWA_WINDOW
         self.is_encoder_decoder = False
         self.is_multimodal = False
+        self.model_is_mrope = False
         self.is_generation = True
         self.quantization = None
         self.is_hybrid_swa = False

python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py

@@ -180,6 +180,7 @@ class TinyGDNModelConfig:
         self.swa_v_head_dim = head_dim
         self.is_encoder_decoder = False
         self.is_multimodal = False
+        self.model_is_mrope = False
         self.is_generation = True
         self.quantization = None
         self.is_hybrid_swa = False

python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py

@@ -183,6 +183,7 @@ class TinyKDAModelConfig:
         self.swa_v_head_dim = head_dim
         self.is_encoder_decoder = False
         self.is_multimodal = False
+        self.model_is_mrope = False
         self.is_generation = True
         self.quantization = None
         self.is_hybrid_swa = False

python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py

@@ -186,6 +186,7 @@ class TinyLightningModelConfig:
         self.swa_v_head_dim = head_dim
         self.is_encoder_decoder = False
         self.is_multimodal = False
+        self.model_is_mrope = False
         self.is_generation = True
         self.quantization = None
         self.is_hybrid_swa = False

python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py

@@ -276,6 +276,7 @@ class TinyMamba2ModelConfig:
         self.swa_v_head_dim = case.head_dim
         self.is_encoder_decoder = False
         self.is_multimodal = False
+        self.model_is_mrope = False
         self.is_generation = True
         self.quantization = None
         self.is_hybrid_swa = False

python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py

@@ -177,6 +177,7 @@ class TinyMLAModelConfig:
         self.scaling = self.head_dim**-0.5
         self.is_encoder_decoder = False
         self.is_multimodal = False
+        self.model_is_mrope = False
         self.is_generation = True
         self.quantization = None
         self.is_hybrid_swa = False