[Gemma 4] Adding MTP support

python/sglang/srt/models/gemma4_causal.py

@@ -878,6 +878,9 @@ def __init__(
     def get_input_embeddings(self) -> nn.Embedding:
         return self.model.embed_tokens
 
+    def get_embed_and_head(self) -> Tuple[torch.Tensor, torch.Tensor]:
+        return self.model.embed_tokens.weight, self.lm_head.weight
+
     def get_attention_sliding_window_size(self):
         return get_attention_sliding_window_size(self.config)
 

python/sglang/srt/models/gemma4_mm.py

@@ -256,6 +256,11 @@ def pad_input_ids(
     def get_input_embeddings(self) -> nn.Embedding:
         return self.language_model.get_input_embeddings()
 
+    def get_embed_and_head(self) -> Tuple[torch.Tensor, torch.Tensor]:
+        # Gemma 4 multimodal ties its LM head to the text embed_tokens
+        embed = self.language_model.embed_tokens.weight
+        return embed, embed
+
     def get_attention_sliding_window_size(self):
         return getattr(self.config.text_config, "sliding_window", -1) - 1