From 89ff2bc1115c5dd33bbb5dba4ccf4f615329b411 Mon Sep 17 00:00:00 2001 From: Charles Chen Date: Sat, 23 May 2026 03:16:47 -0700 Subject: [PATCH] [bug fix] Fix 3 issues when using Gemma4 MTP (#26026) --- python/sglang/srt/models/gemma4_causal.py | 17 +++++++++++------ python/sglang/srt/models/gemma4_mtp.py | 2 ++ python/sglang/srt/server_args.py | 12 +++++++----- 3 files changed, 20 insertions(+), 11 deletions(-) diff --git a/python/sglang/srt/models/gemma4_causal.py b/python/sglang/srt/models/gemma4_causal.py index 190452fcd..c406f12a2 100644 --- a/python/sglang/srt/models/gemma4_causal.py +++ b/python/sglang/srt/models/gemma4_causal.py @@ -1147,7 +1147,8 @@ class Gemma4ForCausalLM(PreTrainedModel): ("experts.w13_weight", "experts.gate_up_proj", ("w1", "w3")), ("experts.w2_weight", "experts.down_proj", ("w2",)), ] - num_experts = self.config.num_experts + # Dense subclasses (e.g. the Gemma4 MTP assistant) reuse this. + num_experts = getattr(self.config, "num_experts", None) or 0 # Per-expert checkpoint format used by compressed-tensors / FP8 # (e.g. RedHatAI/*-FP8-Dynamic) and by ModelOpt NVFP4 @@ -1159,11 +1160,15 @@ class Gemma4ForCausalLM(PreTrainedModel): # in a trailing dot, so the standard `name.replace(weight_name, # param_name)` collapses every suffix uniformly to the fused # FusedMoE params (experts.w13_*, experts.w2_*). - per_expert_params_mapping = FusedMoE.make_expert_params_mapping( - ckpt_gate_proj_name="gate_proj", - ckpt_down_proj_name="down_proj", - ckpt_up_proj_name="up_proj", - num_experts=num_experts, + per_expert_params_mapping = ( + FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=num_experts, + ) + if num_experts + else [] ) k_eq_v_layers = self._get_k_eq_v_layers() diff --git a/python/sglang/srt/models/gemma4_mtp.py b/python/sglang/srt/models/gemma4_mtp.py index 1cb87b7c2..ade10ce5b 100644 --- a/python/sglang/srt/models/gemma4_mtp.py +++ b/python/sglang/srt/models/gemma4_mtp.py @@ -21,6 +21,7 @@ import torch from torch import nn from transformers import PretrainedConfig, PreTrainedModel +from sglang.srt.distributed import get_pp_group from sglang.srt.layers.linear import ReplicatedLinear from sglang.srt.layers.logits_processor import ( LogitsMetadata, @@ -72,6 +73,7 @@ class Gemma4AssistantForCausalLM(Gemma4ForCausalLM): self.assistant_config = config self.config = text_config self.quant_config = quant_config + self.pp_group = get_pp_group() self.vocab_size = text_config.vocab_size self.hidden_size = text_config.hidden_size diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 2b65605d1..2ad0d6226 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -2277,11 +2277,13 @@ class ServerArgs: ) if is_sm100_supported() and self.moe_runner_backend == "auto": - - self.moe_runner_backend = "flashinfer_trtllm" - logger.info( - "Use flashinfer_trtllm as MoE runner backend on SM100 for Gemma-4 NVFP4" - ) + if self.get_model_config().quantization == "modelopt_fp4": + self.quantization = "modelopt_fp4" + self.moe_runner_backend = "flashinfer_trtllm" + logger.info( + "Use flashinfer_trtllm as MoE runner backend on " + "SM100 for Gemma-4 (modelopt_fp4)" + ) elif model_arch == "MossVLForConditionalGeneration": if self.is_attention_backend_not_set(): self.prefill_attention_backend = "flashinfer"