chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Alison Shao <alisonshao@MacBook-Pro-D2W773R9CD.local>
This commit is contained in:
Alison Shao
2026-02-27 14:24:02 -08:00
committed by GitHub
parent e6da514c2c
commit 6ca7da3e7c
3 changed files with 15 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
python/sglang/srt/models/gemma3n_causal.py
View File

@@ -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)

4
python/sglang/srt/models/gemma3n_mm.py
View File

@@ -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,