minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[model] support encoder-free unified Text/Vision/Audio model (#27167)

Browse Source
This commit is contained in:
Xinyuan Tong
2026-06-03 23:58:06 +08:00
committed by GitHub
parent f65aae8493
commit fa5c8a3101
11 changed files with 816 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/speculative_hook.py
+4 -2
View File
@@ -25,8 +25,10 @@ def _resolve_speculative_algorithm_alias(
cfg = get_config(
speculative_draft_model_path, trust_remote_code=trust_remote_code, **kwargs
)
is_gemma4_draft = "Gemma4AssistantForCausalLM" in (
getattr(cfg, "architectures", None) or []
draft_archs = getattr(cfg, "architectures", None) or []
is_gemma4_draft = any(
arch in ("Gemma4AssistantForCausalLM", "Gemma4UnifiedAssistantForCausalLM")
for arch in draft_archs
)
if speculative_algorithm == "EAGLE3" and is_gemma4_draft:
python/sglang/srt/configs/model_config.py
+4
View File
@@ -531,6 +531,7 @@ class ModelConfig:
"MiMoV2MTP",
"Gemma4ForCausalLM",
"Gemma4ForConditionalGeneration",
"Gemma4UnifiedForConditionalGeneration",
]
def _detect_attention_sinks(self) -> bool:
@@ -1515,6 +1516,7 @@ multimodal_model_archs = [
"Gemma3ForConditionalGeneration",
"Gemma3nForConditionalGeneration",
"Gemma4ForConditionalGeneration",
"Gemma4UnifiedForConditionalGeneration",
"Glm4vForConditionalGeneration",
"Glm4vMoeForConditionalGeneration",
"GlmOcrForConditionalGeneration",
@@ -1689,6 +1691,7 @@ def is_hybrid_swa_model(
"Step3p7ForConditionalGeneration",
"Gemma4ForCausalLM",
"Gemma4ForConditionalGeneration",
"Gemma4UnifiedForConditionalGeneration",
"LagunaForCausalLM",
}
if any(arch in hybrid_swa_archs for arch in model_architectures):
@@ -1752,6 +1755,7 @@ def get_hybrid_layer_ids(
elif (
"Gemma4ForCausalLM" in model_architectures
or "Gemma4ForConditionalGeneration" in model_architectures
or "Gemma4UnifiedForConditionalGeneration" in model_architectures
):
layer_types = getattr(hf_text_config, "layer_types", [])
swa_attention_layer_ids = [
python/sglang/srt/entrypoints/openai/serving_chat.py
+2 -1
View File
@@ -174,7 +174,8 @@ class OpenAIServingChat(OpenAIServingBase):
self.is_gemma4 = (
hasattr(self.tokenizer_manager.model_config, "hf_config")
and hasattr(self.tokenizer_manager.model_config.hf_config, "model_type")
and self.tokenizer_manager.model_config.hf_config.model_type == "gemma4"
and self.tokenizer_manager.model_config.hf_config.model_type
in ("gemma4", "gemma4_unified")
)
# Which Python-based chat encoder (if any) bypasses apply_chat_template.
python/sglang/srt/models/gemma4_mtp.py
+5 -1
View File
@@ -397,4 +397,8 @@ class Gemma4AssistantForCausalLM(Gemma4ForCausalLM):
)
EntryClass = Gemma4AssistantForCausalLM
class Gemma4UnifiedAssistantForCausalLM(Gemma4AssistantForCausalLM):
"""Gemma 4 unified MTP assistant; text path identical to the gemma4 assistant."""
EntryClass = [Gemma4AssistantForCausalLM, Gemma4UnifiedAssistantForCausalLM]
python/sglang/srt/models/gemma4_unified.py
+438
View File
@@ -0,0 +1,438 @@
# Copyright 2026 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Gemma4 *Unified* (encoder-free) multimodal model for SGLang.
The unified Gemma4 family (e.g. ``google/gemma-4-12B-it``, arch
``Gemma4UnifiedForConditionalGeneration``, ``model_type="gemma4_unified"``)
shares the Gemma4 *text* decoder verbatim but replaces both modality towers
with light "encoder-free" projection pipelines:
* **Vision** — raw merged pixel patches are projected directly into LM space:
``LN -> Dense -> LN -> +factorized_posemb -> LN`` (``Gemma4UnifiedVisionEmbedder``)
followed by ``RMSNorm -> Linear`` (``Gemma4UnifiedMultimodalEmbedder``).
There is **no** SigLIP attention tower.
* **Audio** — raw 16 kHz waveform is chunked into fixed ``audio_samples_per_token``
frames and projected straight through ``RMSNorm -> Linear``. There is **no**
conformer/USM encoder and no mel spectrogram.
Because the text path is identical to ``gemma4``, we reuse ``Gemma4TextModel``
and subclass ``Gemma4ForConditionalGeneration`` (reusing its ``forward``,
bidirectional-image ``prepare_attn_masks`` and PP/embed plumbing), overriding
only construction, per-modality feature extraction and weight loading.
"""
import logging
import re
from typing import Iterable, List, Optional, Set, Tuple
import torch
from torch import nn
from transformers import PreTrainedModel
from sglang.srt.distributed import get_pp_group
from sglang.srt.layers.layernorm import Gemma4RMSNorm
from sglang.srt.layers.logits_processor import LogitsProcessor, LogitsProcessorOutput
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.utils import PPMissingLayer
from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead
from sglang.srt.managers.schedule_batch import (
MultimodalDataItem,
flatten_nested_list,
)
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.models.gemma4_causal import Gemma4TextModel, pp_filter_load_weight
from sglang.srt.models.gemma4_mm import Gemma4ForConditionalGeneration
from sglang.srt.utils import add_prefix
logger = logging.getLogger(__name__)
class Gemma4UnifiedVisionEmbedder(nn.Module):
"""Encoder-free vision embedder.
Projects raw merged pixel patches ``(..., model_patch_size**2 * 3)`` into
``mm_embed_dim`` via ``LN1 -> Dense -> LN2``, adds factorized 2D positional
embeddings, and applies a final ``LN``. Mirrors HF
``Gemma4UnifiedVisionEmbedder``; runs on the first PP rank only, so it uses
plain (un-sharded) ``nn`` modules.
"""
def __init__(self, config):
super().__init__()
patch_dim = config.model_patch_size**2 * 3 # 48*48*3 = 6912
mm_embed_dim = config.mm_embed_dim
self.patch_ln1 = nn.LayerNorm(patch_dim)
self.patch_dense = nn.Linear(patch_dim, mm_embed_dim)
self.patch_ln2 = nn.LayerNorm(mm_embed_dim)
# Factorized 2D positional embedding table: (mm_posemb_size, 2, mm_embed_dim)
self.pos_embedding = nn.Parameter(
torch.zeros(config.mm_posemb_size, 2, mm_embed_dim)
)
self.pos_norm = nn.LayerNorm(mm_embed_dim)
def forward(
self, pixel_values: torch.Tensor, image_position_ids: torch.Tensor
) -> torch.Tensor:
# pixel_values: (B, num_patches, patch_dim); image_position_ids: (B, num_patches, 2)
hidden_states = self.patch_ln1(pixel_values.to(self.patch_dense.weight.dtype))
hidden_states = self.patch_dense(hidden_states)
hidden_states = self.patch_ln2(hidden_states)
clamped = image_position_ids.clamp(min=0).long()
valid = (image_position_ids != -1).to(self.pos_embedding.dtype).unsqueeze(-1)
axes = torch.arange(2, device=image_position_ids.device)
pos_embs = (self.pos_embedding[clamped, axes] * valid).sum(-2)
hidden_states = hidden_states + pos_embs
hidden_states = self.pos_norm(hidden_states)
return hidden_states
class Gemma4UnifiedMultimodalEmbedder(nn.Module):
"""Shared vision/audio projection: ``RMSNorm(no scale) -> Linear`` to LM space.
Both the vision and audio configs expose ``output_proj_dims`` (the projection
input dim) and ``rms_norm_eps``. ``embedding_pre_projection_norm`` has no
learnable scale, so the only checkpoint tensor is ``embedding_projection.weight``.
"""
def __init__(self, multimodal_config, text_config):
super().__init__()
self.multimodal_hidden_size = multimodal_config.output_proj_dims
self.text_hidden_size = text_config.hidden_size
self.embedding_pre_projection_norm = Gemma4RMSNorm(
self.multimodal_hidden_size,
eps=multimodal_config.rms_norm_eps,
with_scale=False,
)
self.embedding_projection = nn.Linear(
self.multimodal_hidden_size, self.text_hidden_size, bias=False
)
def forward(self, inputs_embeds: torch.Tensor) -> torch.Tensor:
inputs_embeds = inputs_embeds.to(self.embedding_projection.weight.dtype)
normed = self.embedding_pre_projection_norm(inputs_embeds)
return self.embedding_projection(normed)
class Gemma4UnifiedForConditionalGeneration(Gemma4ForConditionalGeneration):
"""Encoder-free unified Gemma4 (text + vision + audio).
Reuses the Gemma4 text decoder and the multimodal ``forward`` / attention
plumbing from :class:`Gemma4ForConditionalGeneration`, swapping the SigLIP
vision tower and conformer audio tower for the encoder-free embedders.
"""
def __init__(
self,
config,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
) -> None:
# Skip Gemma4ForConditionalGeneration.__init__ (it builds the SigLIP /
# conformer towers we do not have) and initialise the HF base directly.
PreTrainedModel.__init__(self, config=config)
self.pp_group = get_pp_group()
self.config = config
self.quant_config = quant_config
text_config = config.text_config
# Encoder-free embedders are consumed only at the input-embedding stage,
# so they live on the first PP rank only.
if self.pp_group.is_first_rank:
self.vision_embedder = (
Gemma4UnifiedVisionEmbedder(config.vision_config)
if getattr(config, "vision_config", None) is not None
else None
)
self.embed_vision = (
Gemma4UnifiedMultimodalEmbedder(config.vision_config, text_config)
if getattr(config, "vision_config", None) is not None
else None
)
self.embed_audio = (
Gemma4UnifiedMultimodalEmbedder(config.audio_config, text_config)
if getattr(config, "audio_config", None) is not None
else None
)
else:
self.vision_embedder = None
self.embed_vision = None
self.embed_audio = None
# Placeholders so methods inherited from the tower-based parent that
# reference these attributes never AttributeError.
self.vision_tower = None
self.audio_tower = None
self.vocab_size = text_config.vocab_size
self.vocab_size_per_layer_input = getattr(
text_config, "vocab_size_per_layer_input", text_config.vocab_size
)
self.language_model = Gemma4TextModel(
text_config,
quant_config,
prefix=add_prefix("language_model", add_prefix("model", prefix)),
)
text_tie = getattr(text_config, "tie_word_embeddings", True)
if self.pp_group.world_size == 1 and text_tie:
self.lm_head = self.language_model.embed_tokens
elif self.pp_group.is_last_rank:
self.lm_head = ParallelLMHead(
text_config.vocab_size,
text_config.hidden_size,
quant_config=quant_config,
prefix=add_prefix("lm_head", prefix),
)
else:
self.lm_head = PPMissingLayer()
self.logits_processor = LogitsProcessor(text_config)
self.capture_aux_hidden_states = False
# The unified checkpoint folds mm-projection vectors into the eoi/eoa
# rows of the (tied) embed_tokens, which inflates their lm-head logits.
# These are input-only markers that must never be sampled — HF applies a
# SuppressTokensLogitsProcessor for exactly these ids
# (generation_config.suppress_tokens). We reproduce that by masking only
# the next-token logits (input-logprob scoring of real eoi/eoa input
# tokens is left untouched).
suppress = []
for attr in ("eoi_token_id", "eoa_token_id", "eoa_token_index"):
tok_id = getattr(config, attr, None)
if isinstance(tok_id, int):
suppress.append(tok_id)
self.suppress_token_ids = sorted(set(suppress))
# Pre-materialize the index as a (non-persistent) buffer so it lives on
# the model's device — indexing with a Python list builds a CPU index
# tensor, which fails CUDA-graph capture ("cannot copy CPU<->CUDA").
self.register_buffer(
"_suppress_idx",
torch.tensor(self.suppress_token_ids, dtype=torch.long),
persistent=False,
)
self.post_init()
@torch.no_grad()
def forward(self, *args, **kwargs):
out = super().forward(*args, **kwargs)
if (
self._suppress_idx.numel() > 0
and isinstance(out, LogitsProcessorOutput)
and out.next_token_logits is not None
):
out.next_token_logits.index_fill_(
1, self._suppress_idx, torch.finfo(out.next_token_logits.dtype).min
)
return out
# ------------------------------------------------------------------
# Per-modality feature extraction (encoder-free)
# ------------------------------------------------------------------
def _empty_embeds(self) -> torch.Tensor:
return torch.empty(
0,
self.language_model.config.hidden_size,
device=next(self.parameters()).device,
dtype=self.language_model.dtype(),
)
def _embed_patches(
self, items: List[MultimodalDataItem], position_attr: str
) -> torch.Tensor:
all_embeds = []
for item in items:
all_pixel_values = flatten_nested_list([item.feature])
all_position_ids = flatten_nested_list([getattr(item, position_attr, None)])
for pv_idx, pv in enumerate(all_pixel_values):
# Pre-embedded passthrough (already at text hidden size).
if (
pv.dim() in (2, 3)
and pv.shape[-1] == self.config.text_config.hidden_size
):
all_embeds.append(pv.to(self.language_model.device))
continue
if pv_idx >= len(all_position_ids) or all_position_ids[pv_idx] is None:
raise ValueError(
f"pixel_values[{pv_idx}] has no matching {position_attr}. "
"The HF image/video processor likely renamed this output — "
"update ATTR_NAME_TO_MODALITY in the Gemma4Unified processor."
)
pp = all_position_ids[pv_idx]
# Collapse video (num_videos, num_frames, P, ...) -> (frames, P, ...)
if pv.dim() == 4:
pv = pv.reshape(-1, pv.shape[-2], pv.shape[-1])
if pp.dim() == 4:
pp = pp.reshape(-1, pp.shape[-2], pp.shape[-1])
if pv.dim() == 2:
pv = pv.unsqueeze(0)
if pp.dim() == 2:
pp = pp.unsqueeze(0)
pv = pv.to(
device=self.language_model.device, dtype=self.language_model.dtype()
)
pp = pp.to(device=self.language_model.device)
embedded = self.vision_embedder(pv, pp) # (B, P, mm_embed_dim)
projected = self.embed_vision(embedded) # (B, P, hidden)
# Drop padding patches (position_ids == -1 on both axes).
padding_mask = (pp == -1).all(dim=-1) # (B, P)
all_embeds.append(projected[~padding_mask])
return torch.cat(all_embeds, dim=0) if all_embeds else self._empty_embeds()
def get_image_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor:
return self._embed_patches(items, "image_position_ids")
def get_video_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor:
return self._embed_patches(items, "video_position_ids")
def get_audio_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor:
if self.embed_audio is None:
raise ValueError(
"Audio inputs provided but the model was built without an audio_config."
)
all_input_features = flatten_nested_list([item.feature for item in items])
# input_features_mask convention: True = valid token.
all_masks = flatten_nested_list([item.input_features_mask for item in items])
all_embeds = []
for input_features, mask in zip(all_input_features, all_masks):
if input_features.dim() == 2:
input_features = input_features.unsqueeze(0)
if mask.dim() == 1:
mask = mask.unsqueeze(0)
input_features = input_features.to(
device=self.language_model.device, dtype=self.language_model.dtype()
)
mask = mask.to(device=input_features.device)
# Raw waveform frames -> RMSNorm -> Linear (no conformer/mel).
projected = self.embed_audio(inputs_embeds=input_features) # (B, T, hidden)
for enc, m in zip(projected, mask):
all_embeds.append(enc[m]) # keep valid frames only
return torch.cat(all_embeds, dim=0) if all_embeds else self._empty_embeds()
# ------------------------------------------------------------------
# Weight loading
# ------------------------------------------------------------------
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> Set[str]:
k_eq_v_layers = self._get_k_eq_v_layers()
params_dict = dict(self.named_parameters())
params_dict.update(dict(self.named_buffers()))
non_persistent_buffers: Set[str] = set()
for mod_name, mod in self.named_modules():
for buf_name in getattr(mod, "_non_persistent_buffers_set", set()):
full = f"{mod_name}.{buf_name}" if mod_name else buf_name
non_persistent_buffers.add(full)
text_tie = getattr(self.config.text_config, "tie_word_embeddings", True)
start_layer = self.language_model.start_layer
end_layer = self.language_model.end_layer
loaded_params: Set[str] = set()
for name, loaded_weight in weights:
name = re.sub(r"^model\.", "", name)
if pp_filter_load_weight(
name,
loaded_weight,
pp_group=self.pp_group,
start_layer=start_layer,
end_layer=end_layer,
params_dict=params_dict,
loaded_params=loaded_params,
tie_word_embeddings=text_tie,
embed_weight_name="language_model.embed_tokens.weight",
first_rank_only_patterns=(
"language_model.embed_tokens",
"language_model.per_layer_model_projection",
"language_model.per_layer_projection_norm",
"vision_embedder.",
"embed_vision.",
"embed_audio.",
),
last_rank_only_prefixes=("language_model.norm.", "lm_head."),
):
continue
# attention_k_eq_v: full-attention layers ship only k_proj (V == K).
# Load k_proj into both the "k" and "v" shards of the fused QKV.
should_dup_k_to_v = (
".k_proj." in name
and k_eq_v_layers
and "language_model." in name
and (m := re.search(r"layers\.(\d+)\.", name)) is not None
and int(m.group(1)) in k_eq_v_layers
)
for param_name, weight_name, shard_id in self.stacked_params_mapping:
if weight_name not in name:
continue
mapped = name.replace(weight_name, param_name)
if mapped not in params_dict:
continue
param = params_dict[mapped]
param.weight_loader(param, loaded_weight, shard_id)
if should_dup_k_to_v:
param.weight_loader(param, loaded_weight, "v")
loaded_params.add(mapped)
break
else:
if name.endswith(".bias") and name not in params_dict:
continue
if name not in params_dict:
continue
param = params_dict[name]
weight_loader = getattr(param, "weight_loader", default_weight_loader)
weight_loader(param, loaded_weight)
loaded_params.add(name)
unloaded_params = params_dict.keys() - loaded_params
if unloaded_params:
param_names = set(dict(self.named_parameters()).keys())
buckets = {
logging.WARNING: (
"Some weights are not initialized from checkpoints",
lambda p: p in param_names,
),
logging.INFO: (
"Persistent buffers not in checkpoint (using default init)",
lambda p: p not in param_names and p not in non_persistent_buffers,
),
logging.DEBUG: (
"Non-persistent buffers not in checkpoint (expected)",
lambda p: p in non_persistent_buffers,
),
}
for level, (msg, pred) in buckets.items():
names = sorted(p for p in unloaded_params if pred(p))
if names:
logger.log(level, "%s: %s", msg, names)
return loaded_params
EntryClass = Gemma4UnifiedForConditionalGeneration
python/sglang/srt/multimodal/processors/base_processor.py
+1
View File
@@ -416,6 +416,7 @@ class BaseMultimodalProcessor(ABC):
if self._processor.__class__.__name__ in {
"Gemma3nProcessor",
"Gemma4Processor",
"Gemma4UnifiedProcessor",
"GlmAsrProcessor",
"Qwen2AudioProcessor",
"Qwen3ASRProcessor",
python/sglang/srt/multimodal/processors/gemma4_unified.py
+33
View File
@@ -0,0 +1,33 @@
# Copyright 2026 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
from sglang.srt.models.gemma4_unified import Gemma4UnifiedForConditionalGeneration
from sglang.srt.multimodal.processors.gemma4 import Gemma4SGLangProcessor
class Gemma4UnifiedSGLangProcessor(Gemma4SGLangProcessor):
"""Multimodal processor for the encoder-free unified Gemma4 (image/video/audio).
Identical to :class:`Gemma4SGLangProcessor` except for audio padding: the
unified model has no SSCP conformer, so the waveform is simply chunked into
fixed ``audio_samples_per_token`` (640) frames. Padding the waveform up to a
multiple of that frame size keeps ``ceil(num_samples / spt)`` consistent with
the number of valid frames the feature extractor emits.
"""
models = [Gemma4UnifiedForConditionalGeneration]
def _get_audio_pad_multiple(self) -> int:
fe = getattr(self._processor, "feature_extractor", None)
return getattr(fe, "audio_samples_per_token", 640)
python/sglang/srt/server_args.py
+1
View File
@@ -2332,6 +2332,7 @@ class ServerArgs:
elif model_arch in (
"Gemma4ForConditionalGeneration",
"Gemma4ForCausalLM",
"Gemma4UnifiedForConditionalGeneration",
):
default_attention_backend = (
"trtllm_mha" if is_sm100_supported() else "triton"
python/sglang/srt/utils/hf_transformers/config.py
+13 -1
View File
@@ -133,7 +133,12 @@ class HfModelConfigParser(ModelConfigParserBase):
if config.model_type == "multi_modality":
_set_architectures(config, "MultiModalityCausalLM")
if config.model_type in ("gemma4", "gemma4_assistant"):
if config.model_type in (
"gemma4",
"gemma4_assistant",
"gemma4_unified",
"gemma4_unified_assistant",
):
# Gemma4 configs use base attributes for SWA layers and `global_*`
# variants for full-attention layers. SGLang expects the opposite:
# base = full-attention, `swa_*` = sliding-window overrides.
@@ -158,6 +163,13 @@ class HfModelConfigParser(ModelConfigParserBase):
if not hasattr(text_config, "swa_v_head_dim"):
text_config.swa_v_head_dim = text_config.swa_head_dim
# Unified Gemma4 names the end-of-audio token `eoa_token_index`,
# but the multimodal processor expects `eoa_token_id`.
if not hasattr(config, "eoa_token_id") and hasattr(
config, "eoa_token_index"
):
config.eoa_token_id = config.eoa_token_index
if config.model_type == "longcat_flash":
_set_architectures(config, "LongcatFlashForCausalLM")