[Model] Support encoder-free unified Text/Vision/Audio model - #27167
Merged
mickqian merged 8 commits intoJun 3, 2026
Conversation
Support the unified Gemma4 family (model_type gemma4_unified, arch Gemma4UnifiedForConditionalGeneration, e.g. gemma-4-12B-it). The unified variant shares the Gemma4 text decoder but replaces both modality towers with encoder-free projection pipelines: vision projects raw merged pixel patches (LN -> Dense -> LN -> +factorized posemb -> LN -> RMSNorm -> Linear) and audio projects raw 16kHz waveform frames (640 samples/token, RMSNorm -> Linear). There is no SigLIP/conformer tower. - models/gemma4_unified.py: Gemma4UnifiedForConditionalGeneration reusing the Gemma4 text decoder; encoder-free vision/audio embedders; weight loading (q/k/v + gate/up fusion, attention_k_eq_v shard duplication). Masks the eoi/eoa next-token logits: the unified checkpoint folds the mm-projection vectors into those tied-embedding rows, inflating their lm-head logits so they get spuriously sampled during text generation; this matches generation_config.suppress_tokens. - multimodal/processors/gemma4_unified.py: processor (640-sample audio framing, no SSCP conv). - utils/hf_transformers/config.py: extend the gemma4 sliding/global head-dim + kv-head remap to gemma4_unified*; map eoa_token_index -> eoa_token_id for the processor. - configs/model_config.py / server_args.py: register the arch (multimodal, hybrid SWA, hybrid layer-id split, default triton/trtllm_mha backend). - multimodal/processors/base_processor.py: route Gemma4UnifiedProcessor audio through the `audio=` kwarg.
Add the unified Text+Vision+Audio gemma-4-12B-it variant to the Gemma 4 cookbook and the deployment-command selector. - Available Models + hardware-requirements tables (1x H200 / 1x B200, TP 1). - Deployment selector: 12B variant for H200 and B200. - Audio Input usage example (the 12B is the audio-capable dense variant). - Speed benchmarks: H200 and B200, text + image, latency + throughput. - Accuracy: MMLU 0.859, MMMU 0.683 (both above the 31B/26B-A4B on MMMU), GSM8K, ASR. GSM8K few-shot under-elicits this reasoning model (0.431); with the chat template it scores 0.950 on the same questions (noted inline). - Config note: on B200/sm100 the default trtllm_mha backend applies causal attention to image tokens; pass --attention-backend triton for multimodal workloads to restore bidirectional image attention. - Perf note: --scheduler-recv-interval 16 lifts B200 text throughput ~3%.
The gemma4 reasoning parser splits on the <|channel>/<channel|> control tokens, which are registered as special tokens, so chat serving must detokenize with skip_special_tokens=False (the existing is_gemma4 gate). The gate matched model_type "gemma4" exactly, so gemma4_unified responses had the markers stripped and the whole completion (answer included) landed in reasoning_content with empty content whenever thinking was enabled. Extend the gate to the gemma4_unified model type.
- Hide the MTP toggle for the 12B in the deployment selector: the paired *-assistant draft uses the gemma4_unified_assistant arch, which SGLang does not register yet, so the generated command would fail at load. - Add the recommended ASR prompt structure to the audio section. - Note that the unified 12B needs a transformers build with the gemma4_unified family (>= 5.10).
Register Gemma4UnifiedAssistantForCausalLM (the gemma4_unified MTP draft; text path identical to the gemma4 assistant) and extend the speculative draft detection so NEXTN/EAGLE is promoted to FROZEN_KV_MTP for it. Restore the MTP toggle for the 12B in the deployment selector and add its MTP launch command to the cookbook. Verified on 1x H200: greedy output identical to non-speculative decoding, thinking/tool-call parsing unaffected, mean accept length ~2.7-3.6 of 6 draft tokens, single-stream decode 109 -> 147 tok/s (~+35%).
JustinTong0323
requested review from
CatherineSue,
ispobock,
kpham-sgl,
merrymercy,
mickqian,
pyc96,
slin1237,
wisclmy0611,
yhyang201,
yuan-luo and
zijiexia
as code owners
June 3, 2026 14:38
Contributor
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
github-actions
Bot
added
documentation
Collaborator
Author
|
/tag-and-rerun-ci |
Collaborator
Author
|
/rerun-failed-ci |
Collaborator
|
please keep an eye on the CI after merging @JustinTong0323 |
2 tasks
kdunn926
added a commit
to kdunn926/OminiX-MLX
that referenced
this pull request
Jun 10, 2026
Port of the merged sglang reference (sgl-project/sglang#27167) for the 'gemma4_unified_vision' schema used by Gemma 4 12B+. Loader + forward only; image preprocessing (HF Gemma4UnifiedImageProcessor equivalent) and API splicing of vision soft tokens into the LM input embedding are deferred to Phase 2. Pipeline (per merged sglang reference): pixel_values[B, P, 48*48*3] -> patch_ln1 (LayerNorm, dim=6912) -> patch_dense (Linear 6912 -> 3840, dequant on load from packed 4-bit triplet) -> patch_ln2 (LayerNorm, dim=3840) -> + factorized 2D pos embed (gather x/y from pos_embedding[N, 2, D], mask -1 padding rows) -> pos_norm (LayerNorm, dim=3840) -> embed_vision (RMSNorm-no-scale + Linear 3840 -> text_hidden, Linear also dequant on load) -> soft tokens[B, P, text_hidden] All checkpoint weights are accounted for: patch_ln1/2/pos_norm weight+bias, patch_dense weight+scales+biases+bias, pos_embedding, embed_vision.embedding_projection weight+scales+biases. Both quantized linears are dequantized to BF16 at load time so the runtime stays on plain nn::Linear and matches the existing canonical EmbedVision path. Smoke test (gemma4-mlx/examples/unified_vision_smoke.rs) loads the 12B-it-4bit checkpoint, runs a forward over synthetic zero-valued patches + a 14x14 position grid, and asserts the output shapes match the reference: pixel_values [1, 196, 6912] -> vision_embedder [1, 196, 3840] -> embed_vision [1, 196, 3840] Pixel-correct end-to-end validation requires the Phase 2 image preprocessor + a real image; the upstream sglang PR was only merged recently so the reference itself is still being shaken down. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Adds SGLang support for the
gemma4_unifiedmodel family(
Gemma4UnifiedForConditionalGeneration) — an encoder-free unifiedText + Vision + Audio architecture. The text path reuses the existing
gemma4decoder; vision/audio use lightweight encoder-free embeddersthat project raw patches / waveform frames directly into the token
stream (no separate vision tower or audio conformer).
Modifications
models/gemma4_unified.py—Gemma4UnifiedForConditionalGeneration(reuses the gemma4 text decoder; adds
Gemma4UnifiedVisionEmbedder+Gemma4UnifiedMultimodalEmbedder; masks the folded mm-projectionlogits on the eoi/eoa rows during next-token sampling).
multimodal/processors/gemma4_unified.py— encoder-free image/audio processor.configs/model_config.py,utils/hf_transformers/config.py— registerthe
gemma4_unified/gemma4_unified_assistantmodel types.entrypoints/openai/serving_chat.py— extend the gemma4 thinking-tokengate to
gemma4_unifiedso the reasoning parser sees the channel tokens.models/gemma4_mtp.py,arg_groups/speculative_hook.py— register theunified MTP assistant; promote
NEXTN/EAGLE→FROZEN_KV_MTPfor theunified draft.
server_args.py,multimodal/processors/base_processor.py— minor wiring.tool-calling, and MTP on H200 and B200 (sm100 uses
--attention-backend tritonfor bidirectional image attention).Accuracy
MMLU / MMMU / GSM8K (chat) match the reference; greedy token-parity vs HF;
audio transcription exact. MTP is lossless (greedy-identical) with ~+35–50%
single-stream decode.
Checklist
CI States
Latest PR Test (Base): ⏳ Run #26892177440
Latest PR Test (Extra): ⏳ Run #26892386220