diff --git a/python/sglang/srt/layers/rotary_embedding/mrope.py b/python/sglang/srt/layers/rotary_embedding/mrope.py index 777edc7c5..b5710d340 100644 --- a/python/sglang/srt/layers/rotary_embedding/mrope.py +++ b/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) diff --git a/python/sglang/srt/managers/mm_utils.py b/python/sglang/srt/managers/mm_utils.py index fdafb44de..e07bfb8c4 100644 --- a/python/sglang/srt/managers/mm_utils.py +++ b/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) diff --git a/python/sglang/srt/models/glm4v.py b/python/sglang/srt/models/glm4v.py index f4df5a2a2..39c7b3caf 100644 --- a/python/sglang/srt/models/glm4v.py +++ b/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 ( diff --git a/python/sglang/srt/models/glm_ocr_nextn.py b/python/sglang/srt/models/glm_ocr_nextn.py index 7e10ae0bd..3f877be60 100644 --- a/python/sglang/srt/models/glm_ocr_nextn.py +++ b/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 diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py index e247d34f8..5851c005f 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/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) diff --git a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py index 4aeb6c5a3..a60549730 100644 --- a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py +++ b/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) diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index e40c64186..16158434f 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/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, diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py index b352cd7c2..9cbbfad9b 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py +++ b/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 diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py index 77bb8e5d4..688c89e72 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py +++ b/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 diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py index a518097aa..065dd4d79 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py +++ b/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 diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py index 3595df14c..fc34825e1 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py +++ b/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 diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py index 4dc11d16c..75787c340 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py +++ b/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 diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py index 471750c7f..b17b663c2 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py +++ b/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 diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py index 46aac0c27..56459d3ef 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py +++ b/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 diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py index 8065fbb8a..e57eb6835 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py +++ b/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 diff --git a/test/registered/unit/spec/test_eagle_draft_cuda_graph_runner.py b/test/registered/unit/spec/test_eagle_draft_cuda_graph_runner.py index 2ab325d8c..25dfe2adc 100644 --- a/test/registered/unit/spec/test_eagle_draft_cuda_graph_runner.py +++ b/test/registered/unit/spec/test_eagle_draft_cuda_graph_runner.py @@ -91,7 +91,7 @@ class TestEagleDraftCudaGraphRunner(CustomTestCase): runner.require_mlp_tp_gather = False runner.require_gathered_buffer = False runner.model_runner = SimpleNamespace( - model_config=SimpleNamespace(vocab_size=8), + model_config=SimpleNamespace(vocab_size=8, model_is_mrope=False), server_args=SimpleNamespace(speculative_use_rejection_sampling=False), device_timer=None, draft_attn_backend=backend, diff --git a/test/registered/unit/spec/test_multimodal_draft_embeddings.py b/test/registered/unit/spec/test_multimodal_draft_embeddings.py new file mode 100644 index 000000000..e51ea27d1 --- /dev/null +++ b/test/registered/unit/spec/test_multimodal_draft_embeddings.py @@ -0,0 +1,59 @@ +import sys +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch + +from sglang.srt.managers.mm_utils import general_mm_embed_routine +from sglang.srt.model_executor.forward_batch_info import ForwardMode +from sglang.srt.runtime_context import get_context +from sglang.srt.speculative.spec_info import SpeculativeAlgorithm +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=2, suite="base-a-test-cpu") + + +class _MutatingLanguageModel: + def get_input_embeddings(self): + return lambda input_ids: torch.zeros((input_ids.shape[0], 2)) + + def __call__(self, *, input_ids, forward_batch, input_embeds): + input_embeds.add_(100) + return input_embeds + + +@pytest.mark.parametrize("algorithm", ["EAGLE", "EAGLE3", "NEXTN", "FROZEN_KV_MTP"]) +def test_eagle_mm_embeddings_survive_in_place_target_updates(algorithm): + source_embeds = torch.arange(8, dtype=torch.float32).reshape(4, 2) + expected_embeds = source_embeds.clone() + forward_batch = SimpleNamespace( + forward_mode=ForwardMode.EXTEND, + contains_mm_inputs=lambda: True, + mm_inputs=[SimpleNamespace()], + extend_prefix_lens_cpu=[0], + extend_seq_lens_cpu=[4], + input_embeds=None, + spec_algorithm=SpeculativeAlgorithm.from_string(algorithm), + ) + + with ( + get_context().override_server_args(), + patch( + "sglang.srt.managers.mm_utils.embed_mm_inputs", + return_value=(source_embeds, {}), + ), + ): + hidden_states = general_mm_embed_routine( + input_ids=torch.arange(4), + forward_batch=forward_batch, + language_model=_MutatingLanguageModel(), + ) + + torch.testing.assert_close(hidden_states, expected_embeds + 100) + torch.testing.assert_close(forward_batch.mm_input_embeds, expected_embeds) + assert forward_batch.mm_input_embeds.data_ptr() != hidden_states.data_ptr() + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__]))