Fix nightly VLM accuracy: gemma3n TP fixes + removal, latency thresholds (#19401)

Co-authored-by: Alison Shao <alisonshao@MacBook-Pro-D2W773R9CD.local>

python/sglang/srt/models/gemma3n_causal.py

@@ -12,6 +12,7 @@ from sglang.srt.layers.linear import (
     ColumnParallelLinear,
     MergedColumnParallelLinear,
     QKVParallelLinear,
+    ReplicatedLinear,
     RowParallelLinear,
 )
 from sglang.srt.layers.logits_processor import LogitsProcessor
@@ -183,21 +184,21 @@ class Gemma3nAltUp(nn.Module):
         self.correct_output_scale = nn.Parameter(
             torch.zeros(config.hidden_size, dtype=torch.float32)
         )
-        self.correction_coefs = ColumnParallelLinear(
+        self.correction_coefs = ReplicatedLinear(
             config.altup_num_inputs,
             config.altup_num_inputs,
             bias=False,
             quant_config=quant_config,
             prefix=add_prefix("correction_coefs", prefix),
         )
-        self.prediction_coefs = ColumnParallelLinear(
+        self.prediction_coefs = ReplicatedLinear(
             config.altup_num_inputs,
             config.altup_num_inputs**2,
             bias=False,
             quant_config=quant_config,
             prefix=add_prefix("prediction_coefs", prefix),
         )
-        self.modality_router = ColumnParallelLinear(
+        self.modality_router = ReplicatedLinear(
             config.hidden_size,
             config.altup_num_inputs,
             bias=False,
@@ -545,14 +546,14 @@ class Gemma3nDecoderLayer(nn.Module):
             config, quant_config, prefix=add_prefix("laurel", prefix)
         )
 
-        self.per_layer_input_gate = ColumnParallelLinear(
+        self.per_layer_input_gate = ReplicatedLinear(
             self.hidden_size,
             self.hidden_size_per_layer_input,
             bias=False,
             quant_config=quant_config,
             prefix=add_prefix("per_layer_input_gate", prefix),
         )
-        self.per_layer_projection = RowParallelLinear(
+        self.per_layer_projection = ReplicatedLinear(
             self.hidden_size_per_layer_input,
             self.hidden_size,
             bias=False,
@@ -677,6 +678,7 @@ class Gemma3nTextModel(PreTrainedModel):
             self.hidden_size,
             config.num_hidden_layers * config.hidden_size_per_layer_input,
             bias=False,
+            gather_output=True,
             quant_config=quant_config,
             prefix=add_prefix("per_layer_model_projection", prefix),
         )
@@ -692,6 +694,7 @@ class Gemma3nTextModel(PreTrainedModel):
                 self.hidden_size,
                 self.hidden_size,
                 bias=False,
+                gather_output=True,
                 quant_config=quant_config,
                 prefix=prefix,
             ),
@@ -704,6 +707,7 @@ class Gemma3nTextModel(PreTrainedModel):
                 self.hidden_size,
                 self.hidden_size,
                 bias=False,
+                gather_output=True,
                 quant_config=quant_config,
                 prefix=prefix,
             ),
@@ -782,9 +786,6 @@ class Gemma3nTextModel(PreTrainedModel):
 
         per_layer_inputs = self.project_per_layer_inputs(input_embeds, per_layer_inputs)
 
-        if positions.dim() == 1:
-            positions = positions.unsqueeze(0)
-
         # Expand hidden_states to support per-layer inputs
         target_magnitude = torch.mean(input_embeds**2, dim=-1, keepdim=True) ** 0.5
         epsilon_tensor = torch.tensor(torch.finfo(input_embeds.dtype).min)

python/sglang/srt/models/gemma3n_mm.py

@@ -14,7 +14,7 @@ from transformers import (
 )
 from transformers.models.auto.modeling_auto import AutoModel
 
-from sglang.srt.layers.linear import RowParallelLinear
+from sglang.srt.layers.linear import ReplicatedLinear
 from sglang.srt.layers.logits_processor import LogitsProcessor
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
 from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
@@ -90,7 +90,7 @@ class Gemma3nMultimodalEmbedder(nn.Module):
             eps=self.eps,
         )
 
-        self.embedding_projection = RowParallelLinear(
+        self.embedding_projection = ReplicatedLinear(
             self.multimodal_hidden_size,
             self.text_hidden_size,
             bias=False,

test/registered/eval/test_vlms_mmmu_eval.py

@@ -34,7 +34,9 @@ MODEL_THRESHOLDS = {
         0.270, 23.8
     ),
     ModelLaunchSettings("google/gemma-3-4b-it"): ModelEvalMetrics(0.360, 10.9),
-    ModelLaunchSettings("google/gemma-3n-E4B-it"): ModelEvalMetrics(0.270, 17.7),
+    ModelLaunchSettings(
+        "google/gemma-3n-E4B-it", extra_args=["--tp=2"]
+    ): ModelEvalMetrics(0.270, 17.7),
     ModelLaunchSettings("mistral-community/pixtral-12b"): ModelEvalMetrics(0.360, 16.6),
     ModelLaunchSettings("moonshotai/Kimi-VL-A3B-Instruct"): ModelEvalMetrics(
         0.330, 23.5
@@ -50,7 +52,7 @@ MODEL_THRESHOLDS = {
     ModelLaunchSettings(
         "unsloth/Mistral-Small-3.1-24B-Instruct-2503"
     ): ModelEvalMetrics(0.30, 16.7),
-    ModelLaunchSettings("XiaomiMiMo/MiMo-VL-7B-RL"): ModelEvalMetrics(0.28, 32.0),
+    ModelLaunchSettings("XiaomiMiMo/MiMo-VL-7B-RL"): ModelEvalMetrics(0.28, 40.0),
     ModelLaunchSettings("zai-org/GLM-4.1V-9B-Thinking"): ModelEvalMetrics(0.280, 30.4),
     ModelLaunchSettings(
         "zai-org/GLM-4.5V-FP8", extra_args=["--tp=2"]