diff --git a/docs_new/cookbook/autoregressive/Google/Gemma4.mdx b/docs_new/cookbook/autoregressive/Google/Gemma4.mdx index a462c2fa5..35a41dba4 100644 --- a/docs_new/cookbook/autoregressive/Google/Gemma4.mdx +++ b/docs_new/cookbook/autoregressive/Google/Gemma4.mdx @@ -47,6 +47,11 @@ Gemma 4 is Google's next-generation family of open models, building on the Gemma Dense ~4B + + [google/gemma-4-12B-it](https://huggingface.co/google/gemma-4-12B-it) + Dense + 12B + [google/gemma-4-31B-it](https://huggingface.co/google/gemma-4-31B-it) Dense @@ -70,6 +75,8 @@ pip install 'git+https://github.com/sgl-project/sglang.git#subdirectory=python' # Install transformers with Gemma 4 support pip install 'git+https://github.com/huggingface/transformers.git@91b1ab1fdfa81a552644a92fbe3e8d88de40e167' +# gemma-4-12B-it (unified) additionally requires a transformers build that +# includes the gemma4_unified model family (>= 5.10). # Or use Docker AMD64 docker pull lmsysorg/sglang:gemma4 # CUDA 12.9 @@ -93,6 +100,7 @@ For the full Docker setup and other installation methods, please refer to the [o ### 3.2 Configuration Tips - SGLang automatically selects the Triton attention backend for Gemma 4 models (required for bidirectional image-token attention during prefill). +- **Attention backend on Blackwell (B200/sm100)**: SGLang defaults to the `trtllm_mha` backend on sm100, which is fastest for text but applies *causal* attention to image tokens. For multimodal (image) workloads on B200, pass `--attention-backend triton` to restore bidirectional image-token attention and full vision quality. Text-only and audio workloads are unaffected by the default. - For the 26B-A4B MoE model, consider `--tp 2` for high-throughput workloads. - **Speculative Decoding (MTP)**: Each Gemma 4 variant ships with a paired `*-assistant` draft model that enables NEXTN multi-token prediction. Enable it via the selector above, or pass `--speculative-algorithm NEXTN --speculative-draft-model-path google/gemma-4--it-assistant --speculative-num-steps 5 --speculative-num-draft-tokens 6 --speculative-eagle-topk 1`. MTP can significantly reduce latency for interactive use cases. The 26B-A4B MoE model requires `--tp 2` when MTP is enabled. - Hardware requirements: @@ -121,6 +129,11 @@ For the full Docker setup and other installation methods, please refer to the [o 1x H200 / 1x MI300X / 1x MI325X / 1x MI355X 1 + + gemma-4-12B-it + 1x H200 / 1x B200 + 1 + gemma-4-31B-it 2x H200 / 1x MI300X / 1x MI325X / 1x MI355X @@ -188,6 +201,18 @@ sglang serve \ --mem-fraction-static 0.85 ``` +```bash Command +# Gemma 4 12B + MTP (~35% faster single-stream decode on H200) +sglang serve \ + --model-path google/gemma-4-12B-it \ + --speculative-algorithm NEXTN \ + --speculative-draft-model-path google/gemma-4-12B-it-assistant \ + --speculative-num-steps 5 \ + --speculative-num-draft-tokens 6 \ + --speculative-eagle-topk 1 \ + --mem-fraction-static 0.85 +``` + ```bash Command # Gemma 4 31B + MTP sglang serve \ @@ -498,6 +523,58 @@ Tool Call: get_weather +### 4.5 Audio Input + +The audio-capable Gemma 4 variants (`gemma-4-E2B-it`, `gemma-4-E4B-it`, `gemma-4-12B-it`) accept raw audio alongside text. Pass the waveform as a base64 `audio_url` data URI (16 kHz mono WAV works well): + +```python Example +import base64 +from openai import OpenAI + +client = OpenAI(base_url="http://localhost:30000/v1", api_key="EMPTY") + +with open("sample.wav", "rb") as f: + audio_b64 = base64.b64encode(f.read()).decode() + +response = client.chat.completions.create( + model="google/gemma-4-12B-it", + messages=[ + { + "role": "user", + "content": [ + {"type": "audio_url", "audio_url": {"url": f"data:audio/wav;base64,{audio_b64}"}}, + {"type": "text", "text": "Transcribe the speech in this audio exactly."}, + ], + } + ], + max_tokens=256, + temperature=0, +) + +print(response.choices[0].message.content) +``` + +
+Example Output + +```text Output +Mr. Quilter is the apostle of the middle classes, and we are glad to welcome his gospel. +``` + +
+ +For best ASR quality, use the recommended transcription prompt structure: + +```text Prompt +Transcribe the following speech segment in {LANGUAGE} into {LANGUAGE} text. + +Follow these specific instructions for formatting the answer: +* Only output the transcription, with no newlines. +* When transcribing numbers, write the digits, i.e. write 1.7 and not one point seven, and write 3 instead of three. +``` + +For speech translation (AST), ask for the transcription in the source language first, then the translation: *"Transcribe the following speech segment in {SOURCE_LANGUAGE}, then translate it into {TARGET_LANGUAGE}. ..."* + ## 5. Benchmark ### 5.1 Speed Benchmark @@ -1033,6 +1110,201 @@ Median ITL (ms): 29.31 ================================================== ``` +#### gemma-4-12B-it (1x H200, TP=1) + +Server Launch Command: +```bash Command +sglang serve --model-path google/gemma-4-12B-it +``` + +**Latency Benchmark (Text)** + +```bash Command +python3 -m sglang.bench_serving --backend sglang \ + --host 0.0.0.0 --port 30000 \ + --dataset-name random --num-prompts 10 --max-concurrency 1 +``` + +```text Output +============ Serving Benchmark Result ============ +Backend: sglang +Max request concurrency: 1 +Successful requests: 10 +Benchmark duration (s): 38.66 +Total input tokens: 6101 +Total generated tokens: 4220 +Request throughput (req/s): 0.26 +Output token throughput (tok/s): 109.15 +Total token throughput (tok/s): 266.94 +Mean TTFT (ms): 33.08 +Median TTFT (ms): 33.71 +Mean TPOT (ms): 9.02 +Median ITL (ms): 9.19 +================================================== +``` + +**Latency Benchmark (Image)** + +```bash Command +python3 -m sglang.bench_serving --backend sglang-oai-chat \ + --host 0.0.0.0 --port 30000 \ + --dataset-name image --image-count 2 --image-resolution 720p \ + --random-input-len 128 --random-output-len 1024 \ + --num-prompts 10 --max-concurrency 1 +``` + +```text Output +============ Serving Benchmark Result ============ +Backend: sglang-oai-chat +Max request concurrency: 1 +Successful requests: 10 +Benchmark duration (s): 39.36 +Total input vision tokens: 5320 +Total generated tokens: 4220 +Request throughput (req/s): 0.25 +Output token throughput (tok/s): 107.23 +Total token throughput (tok/s): 263.62 +Mean TTFT (ms): 94.98 +Median TTFT (ms): 97.33 +Mean TPOT (ms): 9.08 +Median ITL (ms): 9.17 +================================================== +``` + +**Throughput Benchmark (Text)** + +```bash Command +python3 -m sglang.bench_serving --backend sglang \ + --host 0.0.0.0 --port 30000 \ + --dataset-name random --num-prompts 1000 --max-concurrency 100 +``` + +```text Output +============ Serving Benchmark Result ============ +Backend: sglang +Max request concurrency: 100 +Successful requests: 1000 +Benchmark duration (s): 130.44 +Total input tokens: 512842 +Total generated tokens: 510855 +Request throughput (req/s): 7.67 +Output token throughput (tok/s): 3916.46 +Total token throughput (tok/s): 7848.15 +Mean TTFT (ms): 207.49 +Median TTFT (ms): 76.95 +Mean TPOT (ms): 24.38 +Median ITL (ms): 17.89 +================================================== +``` + +**Throughput Benchmark (Image)** + +```text Output +============ Serving Benchmark Result ============ +Backend: sglang-oai-chat +Max request concurrency: 100 +Successful requests: 1000 +Benchmark duration (s): 147.57 +Total input tokens: 619609 +Total input vision tokens: 532000 +Total generated tokens: 510855 +Request throughput (req/s): 6.78 +Output token throughput (tok/s): 3461.79 +Total token throughput (tok/s): 7660.54 +Mean TTFT (ms): 438.40 +Median TTFT (ms): 129.83 +Mean TPOT (ms): 27.12 +Median ITL (ms): 19.16 +================================================== +``` + +#### gemma-4-12B-it (1x B200, TP=1) + +Server Launch Command: +```bash Command +# Text/audio: the sm100 default (trtllm_mha) is fastest. +# For image workloads add --attention-backend triton (bidirectional image attention). +sglang serve --model-path google/gemma-4-12B-it --attention-backend triton +``` + +**Latency Benchmark (Text)** + +```text Output +============ Serving Benchmark Result ============ +Backend: sglang +Max request concurrency: 1 +Successful requests: 10 +Benchmark duration (s): 30.46 +Output token throughput (tok/s): 138.55 +Total token throughput (tok/s): 338.85 +Mean TTFT (ms): 28.14 +Median TTFT (ms): 29.74 +Mean TPOT (ms): 7.08 +Median ITL (ms): 7.26 +================================================== +``` + +**Latency Benchmark (Image)** + +```text Output +============ Serving Benchmark Result ============ +Backend: sglang-oai-chat +Max request concurrency: 1 +Successful requests: 10 +Benchmark duration (s): 31.43 +Total input vision tokens: 5320 +Total generated tokens: 4220 +Request throughput (req/s): 0.32 +Output token throughput (tok/s): 134.26 +Total token throughput (tok/s): 329.57 +Mean TTFT (ms): 115.51 +Median TTFT (ms): 74.27 +Mean TPOT (ms): 7.14 +Median ITL (ms): 7.24 +================================================== +``` + +**Throughput Benchmark (Text)** + +```text Output +============ Serving Benchmark Result ============ +Backend: sglang +Max request concurrency: 100 +Successful requests: 1000 +Benchmark duration (s): 92.94 +Request throughput (req/s): 10.76 +Output token throughput (tok/s): 5496.55 +Total token throughput (tok/s): 11014.49 +Mean TTFT (ms): 120.89 +Median TTFT (ms): 45.00 +Mean TPOT (ms): 17.23 +Median ITL (ms): 14.30 +================================================== +``` + +**Throughput Benchmark (Image)** + +```text Output +============ Serving Benchmark Result ============ +Backend: sglang-oai-chat +Max request concurrency: 100 +Successful requests: 998 +Benchmark duration (s): 107.82 +Total input tokens: 617971 +Total input vision tokens: 530936 +Total generated tokens: 508951 +Request throughput (req/s): 9.26 +Output token throughput (tok/s): 4720.29 +Total token throughput (tok/s): 10451.68 +Mean TTFT (ms): 425.89 +Median TTFT (ms): 109.57 +Mean TPOT (ms): 19.45 +Median ITL (ms): 15.11 +================================================== +``` + +> **Performance tuning**: On B200, raising `--scheduler-recv-interval` to 16 lifted text throughput from 5497 to 5673 tok/s output (≈ +3%) at concurrency 100 with no accuracy change, by reducing the scheduler's per-step Python overhead. It is a safe, low-risk knob for high-concurrency serving. + ### 5.2 Accuracy Benchmark **Test Environment:** @@ -1078,6 +1350,14 @@ Median ITL (ms): 29.31 0.825 **0.810** + + gemma-4-12B-it + 0.784 + 0.888 + 0.946 + 0.861 + **0.859** + gemma-4-31B-it 0.878 @@ -1131,6 +1411,13 @@ Median ITL (ms): 29.31 4.174 4672.030 + + gemma-4-12B-it + 0.431 + 0.052 + 55.105 + 6580.229 + gemma-4-31B-it 0.805 @@ -1148,6 +1435,8 @@ Median ITL (ms): 29.31 +> **Note**: These GSM8K numbers use the raw few-shot completion harness (`sglang.test.few_shot_gsm8k`). `gemma-4-12B-it` is reasoning-oriented and is under-elicited by raw few-shot prompting; with the chat template it scores **0.950** on the same 1319 GSM8K test questions (`sglang.test.run_eval --eval-name gsm8k`). + #### MMMU @@ -1170,6 +1459,10 @@ Median ITL (ms): 29.31 + + + + @@ -1196,6 +1489,12 @@ Median ITL (ms): 29.31 {"Overall-Art and Design": {"num": 120, "acc": 0.458}, "Art": {"num": 30, "acc": 0.433}, "Art_Theory": {"num": 30, "acc": 0.567}, "Design": {"num": 30, "acc": 0.667}, "Music": {"num": 30, "acc": 0.167}, "Overall-Business": {"num": 150, "acc": 0.287}, "Accounting": {"num": 30, "acc": 0.233}, "Economics": {"num": 30, "acc": 0.467}, "Finance": {"num": 30, "acc": 0.133}, "Manage": {"num": 30, "acc": 0.3}, "Marketing": {"num": 30, "acc": 0.3}, "Overall-Science": {"num": 150, "acc": 0.28}, "Biology": {"num": 30, "acc": 0.333}, "Chemistry": {"num": 30, "acc": 0.133}, "Geography": {"num": 30, "acc": 0.4}, "Math": {"num": 30, "acc": 0.2}, "Physics": {"num": 30, "acc": 0.333}, "Overall-Health and Medicine": {"num": 150, "acc": 0.427}, "Basic_Medical_Science": {"num": 30, "acc": 0.4}, "Clinical_Medicine": {"num": 30, "acc": 0.533}, "Diagnostics_and_Laboratory_Medicine": {"num": 30, "acc": 0.4}, "Pharmacy": {"num": 30, "acc": 0.4}, "Public_Health": {"num": 30, "acc": 0.4}, "Overall-Humanities and Social Science": {"num": 120, "acc": 0.7}, "History": {"num": 30, "acc": 0.633}, "Literature": {"num": 30, "acc": 0.867}, "Sociology": {"num": 30, "acc": 0.733}, "Psychology": {"num": 30, "acc": 0.567}, "Overall-Tech and Engineering": {"num": 210, "acc": 0.324}, "Agriculture": {"num": 30, "acc": 0.533}, "Architecture_and_Engineering": {"num": 30, "acc": 0.3}, "Computer_Science": {"num": 30, "acc": 0.367}, "Electronics": {"num": 30, "acc": 0.133}, "Energy_and_Power": {"num": 30, "acc": 0.4}, "Materials": {"num": 30, "acc": 0.2}, "Mechanical_Engineering": {"num": 30, "acc": 0.333}, "Overall": {"num": 900, "acc": 0.396}} ``` +**gemma-4-12B-it** + +```json Config +{"Overall-Art and Design": {"num": 120, "acc": 0.667}, "Art": {"num": 30, "acc": 0.7}, "Art_Theory": {"num": 30, "acc": 0.867}, "Design": {"num": 30, "acc": 0.767}, "Music": {"num": 30, "acc": 0.333}, "Overall-Business": {"num": 150, "acc": 0.747}, "Accounting": {"num": 30, "acc": 0.767}, "Economics": {"num": 30, "acc": 0.767}, "Finance": {"num": 30, "acc": 0.633}, "Manage": {"num": 30, "acc": 0.7}, "Marketing": {"num": 30, "acc": 0.867}, "Overall-Science": {"num": 150, "acc": 0.647}, "Biology": {"num": 30, "acc": 0.6}, "Chemistry": {"num": 30, "acc": 0.633}, "Geography": {"num": 30, "acc": 0.567}, "Math": {"num": 30, "acc": 0.6}, "Physics": {"num": 30, "acc": 0.833}, "Overall-Health and Medicine": {"num": 150, "acc": 0.68}, "Basic_Medical_Science": {"num": 30, "acc": 0.667}, "Clinical_Medicine": {"num": 30, "acc": 0.633}, "Diagnostics_and_Laboratory_Medicine": {"num": 30, "acc": 0.267}, "Pharmacy": {"num": 30, "acc": 0.833}, "Public_Health": {"num": 30, "acc": 1.0}, "Overall-Humanities and Social Science": {"num": 120, "acc": 0.817}, "History": {"num": 30, "acc": 0.8}, "Literature": {"num": 30, "acc": 0.9}, "Sociology": {"num": 30, "acc": 0.8}, "Psychology": {"num": 30, "acc": 0.767}, "Overall-Tech and Engineering": {"num": 210, "acc": 0.6}, "Agriculture": {"num": 30, "acc": 0.467}, "Architecture_and_Engineering": {"num": 30, "acc": 0.667}, "Computer_Science": {"num": 30, "acc": 0.733}, "Electronics": {"num": 30, "acc": 0.567}, "Energy_and_Power": {"num": 30, "acc": 0.667}, "Materials": {"num": 30, "acc": 0.567}, "Mechanical_Engineering": {"num": 30, "acc": 0.533}, "Overall": {"num": 900, "acc": 0.683}} +``` + **gemma-4-31B-it** ```json Config @@ -1238,6 +1537,12 @@ Median ITL (ms): 29.31 + + + + + + @@ -1283,6 +1588,12 @@ Median ITL (ms): 29.31 + + + + + + diff --git a/docs_new/src/snippets/autoregressive/gemma4-deployment.jsx b/docs_new/src/snippets/autoregressive/gemma4-deployment.jsx index 47a907584..36e60e5f9 100644 --- a/docs_new/src/snippets/autoregressive/gemma4-deployment.jsx +++ b/docs_new/src/snippets/autoregressive/gemma4-deployment.jsx @@ -6,6 +6,7 @@ export const Gemma4Deployment = () => { items: [ { id: 'e2b', label: 'E2B (~2B)', default: false }, { id: 'e4b', label: 'E4B (~4B)', default: true }, + { id: '12b', label: '12B (Dense)', default: false }, { id: '31b', label: '31B (Dense)', default: false }, { id: '26b-a4b', label: '26B-A4B (MoE)', default: false }, ] @@ -56,12 +57,14 @@ export const Gemma4Deployment = () => { h200: { e2b: { tp: 1, mem: 0.85 }, e4b: { tp: 1, mem: 0.85 }, + '12b': { tp: 1, mem: 0.85 }, '31b': { tp: 2, mem: 0.85 }, '26b-a4b': { tp: 1, mem: 0.85 }, }, b200: { e2b: { tp: 1, mem: 0.9 }, e4b: { tp: 1, mem: 0.9 }, + '12b': { tp: 1, mem: 0.9 }, '31b': { tp: 1, mem: 0.9 }, '26b-a4b': { tp: 1, mem: 0.9 }, }, @@ -82,6 +85,7 @@ export const Gemma4Deployment = () => { const modelNames = { 'e2b': 'google/gemma-4-E2B-it', 'e4b': 'google/gemma-4-E4B-it', + '12b': 'google/gemma-4-12B-it', '31b': 'google/gemma-4-31B-it', '26b-a4b': 'google/gemma-4-26B-A4B-it', }; diff --git a/python/sglang/srt/arg_groups/speculative_hook.py b/python/sglang/srt/arg_groups/speculative_hook.py index 1c38fcb40..71008d6cd 100644 --- a/python/sglang/srt/arg_groups/speculative_hook.py +++ b/python/sglang/srt/arg_groups/speculative_hook.py @@ -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: diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 929adc264..f01ec9f6e 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -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 = [ diff --git a/python/sglang/srt/entrypoints/openai/serving_chat.py b/python/sglang/srt/entrypoints/openai/serving_chat.py index cc02ebd6a..cbd3df776 100644 --- a/python/sglang/srt/entrypoints/openai/serving_chat.py +++ b/python/sglang/srt/entrypoints/openai/serving_chat.py @@ -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. diff --git a/python/sglang/srt/models/gemma4_mtp.py b/python/sglang/srt/models/gemma4_mtp.py index ade10ce5b..f652661fc 100644 --- a/python/sglang/srt/models/gemma4_mtp.py +++ b/python/sglang/srt/models/gemma4_mtp.py @@ -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] diff --git a/python/sglang/srt/models/gemma4_unified.py b/python/sglang/srt/models/gemma4_unified.py new file mode 100644 index 000000000..67d0771d4 --- /dev/null +++ b/python/sglang/srt/models/gemma4_unified.py @@ -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 diff --git a/python/sglang/srt/multimodal/processors/base_processor.py b/python/sglang/srt/multimodal/processors/base_processor.py index 3dca4f28d..964583a0e 100644 --- a/python/sglang/srt/multimodal/processors/base_processor.py +++ b/python/sglang/srt/multimodal/processors/base_processor.py @@ -416,6 +416,7 @@ class BaseMultimodalProcessor(ABC): if self._processor.__class__.__name__ in { "Gemma3nProcessor", "Gemma4Processor", + "Gemma4UnifiedProcessor", "GlmAsrProcessor", "Qwen2AudioProcessor", "Qwen3ASRProcessor", diff --git a/python/sglang/srt/multimodal/processors/gemma4_unified.py b/python/sglang/srt/multimodal/processors/gemma4_unified.py new file mode 100644 index 000000000..676f11f34 --- /dev/null +++ b/python/sglang/srt/multimodal/processors/gemma4_unified.py @@ -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) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 5cb20a067..2e4b72d8a 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -2332,6 +2332,7 @@ class ServerArgs: elif model_arch in ( "Gemma4ForConditionalGeneration", "Gemma4ForCausalLM", + "Gemma4UnifiedForConditionalGeneration", ): default_attention_backend = ( "trtllm_mha" if is_sm100_supported() else "triton" diff --git a/python/sglang/srt/utils/hf_transformers/config.py b/python/sglang/srt/utils/hf_transformers/config.py index a46c4a45d..f85a25f94 100644 --- a/python/sglang/srt/utils/hf_transformers/config.py +++ b/python/sglang/srt/utils/hf_transformers/config.py @@ -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")
gemma-4-E4B-it **0.396**
gemma-4-12B-it**0.683**
gemma-4-31B-it **0.589**0.366 2.46
gemma-4-12B-itSupported (see §4.5)——
gemma-4-31B-it Not Supported0.8707s 16.20
gemma-4-12B-itSupported (see §4.5)——
gemma-4-31B-it Not Supported