[Feature] Support Gemma3-Vision

[pengchengneo]

Checklist

• 1. If the issue you raised is not a feature but a question, please raise a discussion at https://github.com/sgl-project/sgl-jax/discussions/new/choose Otherwise, it will be closed.
• 2. Please use English, otherwise it will be closed.

Motivation

Support new multimodal modal Gemma3-Vision

Related resources

as title

[jackopenn]

Hey, I'd like to work on this. :)

[pengchengneo]

Hey, I'd like to work on this. :)

Great. Before you begin, may I ask if you have sufficient TPU resources for development and testing? Additionally, please provide an RFC to ensure your architecture aligns with our design.

[jackopenn]

Great. Before you begin, may I ask if you have sufficient TPU resources for development and testing?

Yes I do, although I will most likely request access to resources (skypilot) via slack.

Additionally, please provide an RFC to ensure your architecture aligns with our design.

Will do :)

[jackopenn]

[RFC] Support Gemma3-Vision in SGL_JAX

Summary

Add Gemma 3 to SGLang-JAX: the 1B text-only model and the 4B, 12B, 27B vision variants. The vision models use the existing 2-stage multimodal pipeline (#552), combining a frozen SigLIP vision encoder with a Gemma 3 text decoder to produce a fixed 256 vision tokens per image.

This PR leverages the existing multimodal strategy, following a similar approach to how Qwen 2.5 VL is implemented. The main work is implementing the new model modules and wiring them into the registry.

What This PR Adds

New Files

File	Description
srt/models/gemma3.py	Gemma3 text decoder — reuses Gemma2 layers, adds 5:1 local/global attention pattern
srt/multimodal/models/gemma3/gemma3_vision_encoder.py	SigLIP vision encoder (Conv2D patch embed + learned pos embed + transformer) and multimodal projector (AvgPool → RMSNorm → Linear)
srt/multimodal/models/gemma3/gemma3_decoder.py	Gemma3 decoder wrapper for Stage 1
srt/multimodal/configs/gemma3/gemma3_vision_config.py	Gemma3VisionConfig dataclass
srt/multimodal/models/static_configs/gemma3_stage_config.yaml	2-stage pipeline config (TP=1)
srt/multimodal/models/static_configs/gemma3_stage_config_tp4.yaml	2-stage pipeline config (TP=4, for 27B)

Modified Files

File	Change
srt/multimodal/models/static_configs/yaml_registry.py	Register gemma-3-* model names → stage config YAML
srt/multimodal/configs/config_registry.py	Add Gemma3VisionConfigRegistry
srt/multimodal/manager/stage.py	Register Gemma3VisionModel and Gemma3Decoder in get_model_class()
srt/multimodal/model_executor/vit/vit_model_runner.py	Dispatch between Qwen ViT and SigLIP — Gemma3 needs no aux arrays
srt/multimodal/manager/multimodal_tokenizer.py	Handle Gemma3 config attribute names (boi_token_index/eoi_token_index)

Key Design Decisions

model_config.py already has Gemma3 registered

Gemma3ForConditionalGeneration is already in multimodal_model_archs (line 681 of srt/configs/model_config.py). The text-only 1B uses Gemma3ForCausalLM which auto-routes to the text model registry. No changes needed to model_config.py.

Architectural differences from Gemma2 (for srt/models/gemma3.py)

The text decoder reuses Gemma2 layers, but with these differences:

• No logit softcapping — Gemma2 has final_logit_softcapping and attn_logit_softcapping; Gemma3 drops both
• GELU activation — Gemma2 uses GELU-tanh approximation; Gemma3 uses standard GELU
• 5:1 sliding window — config.layer_types alternates "sliding_attention" and "global_attention" at a 5:1 ratio (Gemma2 alternates 1:1)
• Different head dimensions — Gemma3 1B has head_dim=256 (vs typical 128)

Weight loading

Gemma3 text decoder weight mappings are structurally identical to Gemma2 — same layer prefix pattern, same projection names. The _create_layer_mappings can be reused nearly verbatim, just without softcapping-related params.

Testing Strategy

Unit tests (python/sgl_jax/test/):

• Vision encoder output shapes (pixel_values → hidden states)
• Projector output shapes (4096 patches → 256 tokens)
• Weight mapping coverage (HF checkpoint keys → NNX parameters)

E2E test (test/srt/):

• Launch server with google/gemma-3-4b-it
• Send image chat completion via OpenAI-compatible API
• Verify non-empty, coherent response

CI integration:

• Unit tests added to unit-test-tpu-v6e-1 suite
• E2E test added to per-commit or nightly suite

Milestones

• Milestone 1: Gemma3 text-only model — text decoder, weight loading, tested on 1B
  • Text model loads and produces logits (weight mapping coverage)
  • E2E: launch server with google/gemma-3-1b-it, send text chat completion, verify response
• Milestone 2: Vision support for 4B/12B/27B — SigLIP encoder, projector, pipeline wiring, unit + E2E tests
• Milestone 3: TP=4 for 27B, Pan & Scan for high-res images, performance benchmarking

References

• Technical Report: https://arxiv.org/abs/2503.19786
• Model Weights: 1B / 4B / 12B / 27B
• HuggingFace Docs: https://huggingface.co/docs/transformers/en/model_doc/gemma3
• SGLang Gemma3 PR: model: support gemma-3-it sglang#4424
• SGLang-JAX multimodal RFC: [RFC]: Support multimodal model in SGL_JAX #552

[pengchengneo]

Well done, please get right in @jackopenn

[aolemila]

In order to integrate with Tunix, we need to add mapping logic in Tunix. @pengchengneo